library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub maspypy/library

:heavy_check_mark: poly/ntt_doubling.hpp

Depends on

Required by

Verified with

Code

#include "poly/ntt.hpp"

// 2^k 次多項式の長さ 2^k が与えられるので 2^k+1 にする
template <typename mint, bool transposed = false>
void ntt_doubling(vector<mint>& a) {
  static array<mint, 30> root;
  static bool prepared = 0;
  if (!prepared) {
    prepared = 1;
    const int rank2 = mint::ntt_info().fi;
    root[rank2] = mint::ntt_info().se;
    FOR_R(i, rank2) { root[i] = root[i + 1] * root[i + 1]; }
  }

  if constexpr (!transposed) {
    const int M = (int)a.size();
    auto b = a;
    ntt(b, 1);
    mint r = 1, zeta = root[topbit(2 * M)];
    FOR(i, M) b[i] *= r, r *= zeta;
    ntt(b, 0);
    copy(begin(b), end(b), back_inserter(a));
  } else {
    const int M = len(a) / 2;
    vc<mint> tmp = {a.begin(), a.begin() + M};
    a = {a.begin() + M, a.end()};
    transposed_ntt(a, 0);
    mint r = 1, zeta = root[topbit(2 * M)];
    FOR(i, M) a[i] *= r, r *= zeta;
    transposed_ntt(a, 1);
    FOR(i, M) a[i] += tmp[i];
  }
}
#line 1 "poly/ntt_doubling.hpp"

#line 1 "other/bit.hpp"

int popcnt(int x) { return __builtin_popcount(x); }
int popcnt(u32 x) { return __builtin_popcount(x); }
int popcnt(ll x) { return __builtin_popcountll(x); }
int popcnt(u64 x) { return __builtin_popcountll(x); }
int popcnt_sgn(int x) { return (__builtin_parity(unsigned(x)) & 1 ? -1 : 1); }
int popcnt_sgn(u32 x) { return (__builtin_parity(x) & 1 ? -1 : 1); }
int popcnt_sgn(ll x) { return (__builtin_parityll(x) & 1 ? -1 : 1); }
int popcnt_sgn(u64 x) { return (__builtin_parityll(x) & 1 ? -1 : 1); }
// (0, 1, 2, 3, 4) -> (-1, 0, 1, 1, 2)
int topbit(int x) { return (x == 0 ? -1 : 31 - __builtin_clz(x)); }
int topbit(u32 x) { return (x == 0 ? -1 : 31 - __builtin_clz(x)); }
int topbit(ll x) { return (x == 0 ? -1 : 63 - __builtin_clzll(x)); }
int topbit(u64 x) { return (x == 0 ? -1 : 63 - __builtin_clzll(x)); }
// (0, 1, 2, 3, 4) -> (-1, 0, 1, 0, 2)
int lowbit(int x) { return (x == 0 ? -1 : __builtin_ctz(x)); }
int lowbit(u32 x) { return (x == 0 ? -1 : __builtin_ctz(x)); }
int lowbit(ll x) { return (x == 0 ? -1 : __builtin_ctzll(x)); }
int lowbit(u64 x) { return (x == 0 ? -1 : __builtin_ctzll(x)); }

template <typename T>
T kth_bit(int k) {
  return T(1) << k;
}
template <typename T>
bool has_kth_bit(T x, int k) {
  return x >> k & 1;
}

template <typename UINT>
struct all_bit {
  UINT s;
  all_bit(UINT s) : s(s) {}
  struct iter {
    UINT s;
    int operator*() const { return lowbit(s); }
    void operator++() { s &= s - 1; }
    bool operator!=(nullptr_t) const { return s; }
  };
  iter begin() const { return {s}; }
  nullptr_t end() const { return nullptr; }
};

template <typename UINT>
struct all_subset {
  UINT s;
  all_subset(UINT s) : s(s) {}
  struct iter {
    UINT s, t;
    bool done = false;
    UINT operator*() const { return t; }
    void operator++() {
      done = (t == 0);
      t = (t - 1) & s;
    }
    bool operator!=(nullptr_t) const { return !done; }
  };
  iter begin() const { return {s, s}; }
  nullptr_t end() const { return nullptr; }
};

constexpr u64 full_mask(int n) { return n == 64 ? -1ULL : (1ULL << n) - 1; }

u64 bit_reverse(u64 x) {
  x = ((x & 0x5555555555555555ULL) << 1) | ((x >> 1) & 0x5555555555555555ULL);
  x = ((x & 0x3333333333333333ULL) << 2) | ((x >> 2) & 0x3333333333333333ULL);
  x = ((x & 0x0f0f0f0f0f0f0f0fULL) << 4) | ((x >> 4) & 0x0f0f0f0f0f0f0f0fULL);
  x = ((x & 0x00ff00ff00ff00ffULL) << 8) | ((x >> 8) & 0x00ff00ff00ff00ffULL);
  x = ((x & 0x0000ffff0000ffffULL) << 16) | ((x >> 16) & 0x0000ffff0000ffffULL);
  x = (x << 32) | (x >> 32);
  return x;
}
#line 2 "poly/ntt.hpp"

template <class mint>
void ntt(vector<mint>& a, bool inverse) {
  assert(mint::can_ntt());
  const int rank2 = mint::ntt_info().fi;
  const u32 mod = mint::get_mod();
  static array<mint, 30> root, iroot;
  static array<mint, 30> rate2, irate2;
  static array<mint, 30> rate3, irate3;

  assert(rank2 != -1 && len(a) <= (1 << max(0, rank2)));

  static bool prepared = 0;
  if (!prepared) {
    prepared = 1;
    root[rank2] = mint::ntt_info().se;
    iroot[rank2] = mint(1) / root[rank2];
    FOR_R(i, rank2) {
      root[i] = root[i + 1] * root[i + 1];
      iroot[i] = iroot[i + 1] * iroot[i + 1];
    }
    mint prod = 1, iprod = 1;
    for (int i = 0; i <= rank2 - 2; i++) {
      rate2[i] = root[i + 2] * prod;
      irate2[i] = iroot[i + 2] * iprod;
      prod *= iroot[i + 2];
      iprod *= root[i + 2];
    }
    prod = 1, iprod = 1;
    for (int i = 0; i <= rank2 - 3; i++) {
      rate3[i] = root[i + 3] * prod;
      irate3[i] = iroot[i + 3] * iprod;
      prod *= iroot[i + 3];
      iprod *= root[i + 3];
    }
  }

  int n = int(a.size());
  int h = topbit(n);
  assert(n == 1 << h);
  if (!inverse) {
    int len = 0;
    while (len < h) {
      if (h - len == 1) {
        int p = 1 << (h - len - 1);
        mint rot = 1;
        FOR(s, 1 << len) {
          int offset = s << (h - len);
          FOR(i, p) {
            auto l = a[i + offset];
            auto r = a[i + offset + p] * rot;
            a[i + offset] = l + r;
            a[i + offset + p] = l - r;
          }
          rot *= rate2[topbit(~s & -~s)];
        }
        len++;
      } else {
        int p = 1 << (h - len - 2);
        mint rot = 1, imag = root[2];
        for (int s = 0; s < (1 << len); s++) {
          mint rot2 = rot * rot;
          mint rot3 = rot2 * rot;
          int offset = s << (h - len);
          for (int i = 0; i < p; i++) {
            u64 mod2 = u64(mod) * mod;
            u64 a0 = a[i + offset].val;
            u64 a1 = u64(a[i + offset + p].val) * rot.val;
            u64 a2 = u64(a[i + offset + 2 * p].val) * rot2.val;
            u64 a3 = u64(a[i + offset + 3 * p].val) * rot3.val;
            u64 a1na3imag = (a1 + mod2 - a3) % mod * imag.val;
            u64 na2 = mod2 - a2;
            a[i + offset] = a0 + a2 + a1 + a3;
            a[i + offset + 1 * p] = a0 + a2 + (2 * mod2 - (a1 + a3));
            a[i + offset + 2 * p] = a0 + na2 + a1na3imag;
            a[i + offset + 3 * p] = a0 + na2 + (mod2 - a1na3imag);
          }
          rot *= rate3[topbit(~s & -~s)];
        }
        len += 2;
      }
    }
  } else {
    mint coef = mint(1) / mint(len(a));
    FOR(i, len(a)) a[i] *= coef;
    int len = h;
    while (len) {
      if (len == 1) {
        int p = 1 << (h - len);
        mint irot = 1;
        FOR(s, 1 << (len - 1)) {
          int offset = s << (h - len + 1);
          FOR(i, p) {
            u64 l = a[i + offset].val;
            u64 r = a[i + offset + p].val;
            a[i + offset] = l + r;
            a[i + offset + p] = (mod + l - r) * irot.val;
          }
          irot *= irate2[topbit(~s & -~s)];
        }
        len--;
      } else {
        int p = 1 << (h - len);
        mint irot = 1, iimag = iroot[2];
        FOR(s, (1 << (len - 2))) {
          mint irot2 = irot * irot;
          mint irot3 = irot2 * irot;
          int offset = s << (h - len + 2);
          for (int i = 0; i < p; i++) {
            u64 a0 = a[i + offset + 0 * p].val;
            u64 a1 = a[i + offset + 1 * p].val;
            u64 a2 = a[i + offset + 2 * p].val;
            u64 a3 = a[i + offset + 3 * p].val;
            u64 x = (mod + a2 - a3) * iimag.val % mod;
            a[i + offset] = a0 + a1 + a2 + a3;
            a[i + offset + 1 * p] = (a0 + mod - a1 + x) * irot.val;
            a[i + offset + 2 * p] = (a0 + a1 + 2 * mod - a2 - a3) * irot2.val;
            a[i + offset + 3 * p] = (a0 + 2 * mod - a1 - x) * irot3.val;
          }
          irot *= irate3[topbit(~s & -~s)];
        }
        len -= 2;
      }
    }
  }
}
#line 3 "poly/ntt_doubling.hpp"

// 2^k 次多項式の長さ 2^k が与えられるので 2^k+1 にする
template <typename mint, bool transposed = false>
void ntt_doubling(vector<mint>& a) {
  static array<mint, 30> root;
  static bool prepared = 0;
  if (!prepared) {
    prepared = 1;
    const int rank2 = mint::ntt_info().fi;
    root[rank2] = mint::ntt_info().se;
    FOR_R(i, rank2) { root[i] = root[i + 1] * root[i + 1]; }
  }

  if constexpr (!transposed) {
    const int M = (int)a.size();
    auto b = a;
    ntt(b, 1);
    mint r = 1, zeta = root[topbit(2 * M)];
    FOR(i, M) b[i] *= r, r *= zeta;
    ntt(b, 0);
    copy(begin(b), end(b), back_inserter(a));
  } else {
    const int M = len(a) / 2;
    vc<mint> tmp = {a.begin(), a.begin() + M};
    a = {a.begin() + M, a.end()};
    transposed_ntt(a, 0);
    mint r = 1, zeta = root[topbit(2 * M)];
    FOR(i, M) a[i] *= r, r *= zeta;
    transposed_ntt(a, 1);
    FOR(i, M) a[i] += tmp[i];
  }
}
Back to top page