Number Theoretic Transform Friendly ModInt (math/fft/number-theoretic-transform-friendly-mod-int.hpp)

• View this file on GitHub
• Last update: 2024-05-01 23:44:47+09:00
• Include: #include "math/fft/number-theoretic-transform-friendly-mod-int.hpp"

Required by

• Formal Power Series Friendly NTT(NTTmod用形式的冪級数) (math/fps/formal-power-series-friendly-ntt.hpp)

Verified with

• test/verify/yosupo-bernoulli-number.test.cpp
• test/verify/yosupo-convolution-mod-2.test.cpp
• test/verify/yosupo-division-of-polynomials.test.cpp
• test/verify/yosupo-exp-of-formal-power-series.test.cpp
• test/verify/yosupo-find-linear-recurrence.test.cpp
• test/verify/yosupo-inv-of-formal-power-series.test.cpp
• test/verify/yosupo-kth-term-of-linearly-recurrent-sequence.test.cpp
• test/verify/yosupo-log-of-formal-power-series.test.cpp
• test/verify/yosupo-multipoint-evaluation.test.cpp
• test/verify/yosupo-partition-function.test.cpp
• test/verify/yosupo-polynomial-interpolation.test.cpp
• test/verify/yosupo-polynomial-taylor-shift.test.cpp
• test/verify/yosupo-pow-of-formal-power-series.test.cpp
• test/verify/yosupo-sharp-p-subset-sum.test.cpp
• test/verify/yosupo-shift-of-sampling-points-of-polynomial.test.cpp
• test/verify/yosupo-sqrt-of-formal-power-series.test.cpp
• test/verify/yosupo-stirling-number-of-the-first-kind.test.cpp
• test/verify/yosupo-stirling-number-of-the-second-kind.test.cpp

Code

/**
 * @brief Number Theoretic Transform Friendly ModInt
 */
template <typename Mint>
struct NumberTheoreticTransformFriendlyModInt {
  static vector<Mint> roots, iroots, rate3, irate3;
  static int max_base;

  NumberTheoreticTransformFriendlyModInt() = default;

  static void init() {
    if (roots.empty()) {
      const unsigned mod = Mint::mod();
      assert(mod >= 3 && mod % 2 == 1);
      auto tmp = mod - 1;
      max_base = 0;
      while (tmp % 2 == 0) tmp >>= 1, max_base++;
      Mint root = 2;
      while (root.pow((mod - 1) >> 1) == 1) {
        root += 1;
      }
      assert(root.pow(mod - 1) == 1);

      roots.resize(max_base + 1);
      iroots.resize(max_base + 1);
      rate3.resize(max_base + 1);
      irate3.resize(max_base + 1);

      roots[max_base] = root.pow((mod - 1) >> max_base);
      iroots[max_base] = Mint(1) / roots[max_base];
      for (int i = max_base - 1; i >= 0; i--) {
        roots[i] = roots[i + 1] * roots[i + 1];
        iroots[i] = iroots[i + 1] * iroots[i + 1];
      }
      {
        Mint prod = 1, iprod = 1;
        for (int i = 0; i <= max_base - 3; i++) {
          rate3[i] = roots[i + 3] * prod;
          irate3[i] = iroots[i + 3] * iprod;
          prod *= iroots[i + 3];
          iprod *= roots[i + 3];
        }
      }
    }
  }

  static void ntt(vector<Mint> &a) {
    init();
    const int n = (int)a.size();
    assert((n & (n - 1)) == 0);
    int h = __builtin_ctz(n);
    assert(h <= max_base);
    int len = 0;
    Mint imag = roots[2];
    if (h & 1) {
      int p = 1 << (h - 1);
      Mint rot = 1;
      for (int i = 0; i < p; i++) {
        auto r = a[i + p];
        a[i + p] = a[i] - r;
        a[i] += r;
      }
      len++;
    }
    for (; len + 1 < h; len += 2) {
      int p = 1 << (h - len - 2);
      {  // s = 0
        for (int i = 0; i < p; i++) {
          auto a0 = a[i];
          auto a1 = a[i + p];
          auto a2 = a[i + 2 * p];
          auto a3 = a[i + 3 * p];
          auto a1na3imag = (a1 - a3) * imag;
          auto a0a2 = a0 + a2;
          auto a1a3 = a1 + a3;
          auto a0na2 = a0 - a2;
          a[i] = a0a2 + a1a3;
          a[i + 1 * p] = a0a2 - a1a3;
          a[i + 2 * p] = a0na2 + a1na3imag;
          a[i + 3 * p] = a0na2 - a1na3imag;
        }
      }
      Mint rot = rate3[0];
      for (int s = 1; s < (1 << len); s++) {
        int offset = s << (h - len);
        Mint rot2 = rot * rot;
        Mint rot3 = rot2 * rot;
        for (int i = 0; i < p; i++) {
          auto a0 = a[i + offset];
          auto a1 = a[i + offset + p] * rot;
          auto a2 = a[i + offset + 2 * p] * rot2;
          auto a3 = a[i + offset + 3 * p] * rot3;
          auto a1na3imag = (a1 - a3) * imag;
          auto a0a2 = a0 + a2;
          auto a1a3 = a1 + a3;
          auto a0na2 = a0 - a2;
          a[i + offset] = a0a2 + a1a3;
          a[i + offset + 1 * p] = a0a2 - a1a3;
          a[i + offset + 2 * p] = a0na2 + a1na3imag;
          a[i + offset + 3 * p] = a0na2 - a1na3imag;
        }
        rot *= rate3[__builtin_ctz(~s)];
      }
    }
  }

  static void intt(vector<Mint> &a, bool f = true) {
    init();
    const int n = (int)a.size();
    assert((n & (n - 1)) == 0);
    int h = __builtin_ctz(n);
    assert(h <= max_base);
    int len = h;
    Mint iimag = iroots[2];
    for (; len > 1; len -= 2) {
      int p = 1 << (h - len);
      {  // s = 0
        for (int i = 0; i < p; i++) {
          auto a0 = a[i];
          auto a1 = a[i + 1 * p];
          auto a2 = a[i + 2 * p];
          auto a3 = a[i + 3 * p];
          auto a2na3iimag = (a2 - a3) * iimag;
          auto a0na1 = a0 - a1;
          auto a0a1 = a0 + a1;
          auto a2a3 = a2 + a3;
          a[i] = a0a1 + a2a3;
          a[i + 1 * p] = (a0na1 + a2na3iimag);
          a[i + 2 * p] = (a0a1 - a2a3);
          a[i + 3 * p] = (a0na1 - a2na3iimag);
        }
      }
      Mint irot = irate3[0];
      for (int s = 1; s < (1 << (len - 2)); s++) {
        int offset = s << (h - len + 2);
        Mint irot2 = irot * irot;
        Mint irot3 = irot2 * irot;
        for (int i = 0; i < p; i++) {
          auto a0 = a[i + offset];
          auto a1 = a[i + offset + 1 * p];
          auto a2 = a[i + offset + 2 * p];
          auto a3 = a[i + offset + 3 * p];
          auto a2na3iimag = (a2 - a3) * iimag;
          auto a0na1 = a0 - a1;
          auto a0a1 = a0 + a1;
          auto a2a3 = a2 + a3;
          a[i + offset] = a0a1 + a2a3;
          a[i + offset + 1 * p] = (a0na1 + a2na3iimag) * irot;
          a[i + offset + 2 * p] = (a0a1 - a2a3) * irot2;
          a[i + offset + 3 * p] = (a0na1 - a2na3iimag) * irot3;
        }
        irot *= irate3[__builtin_ctz(~s)];
      }
    }
    if (len >= 1) {
      int p = 1 << (h - 1);
      for (int i = 0; i < p; i++) {
        auto ajp = a[i] - a[i + p];
        a[i] += a[i + p];
        a[i + p] = ajp;
      }
    }
    if (f) {
      Mint inv_sz = Mint(1) / n;
      for (int i = 0; i < n; i++) a[i] *= inv_sz;
    }
  }

  static vector<Mint> multiply(vector<Mint> a, vector<Mint> b) {
    int need = a.size() + b.size() - 1;
    int nbase = 1;
    while ((1 << nbase) < need) nbase++;
    int sz = 1 << nbase;
    a.resize(sz, 0);
    b.resize(sz, 0);
    ntt(a);
    ntt(b);
    Mint inv_sz = Mint(1) / sz;
    for (int i = 0; i < sz; i++) a[i] *= b[i] * inv_sz;
    intt(a, false);
    a.resize(need);
    return a;
  }
};

template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::roots =
    vector<Mint>();
template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::iroots =
    vector<Mint>();
template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::rate3 =
    vector<Mint>();
template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::irate3 =
    vector<Mint>();
template <typename Mint>
int NumberTheoreticTransformFriendlyModInt<Mint>::max_base = 0;

#line 1 "math/fft/number-theoretic-transform-friendly-mod-int.hpp"
/**
 * @brief Number Theoretic Transform Friendly ModInt
 */
template <typename Mint>
struct NumberTheoreticTransformFriendlyModInt {
  static vector<Mint> roots, iroots, rate3, irate3;
  static int max_base;

  NumberTheoreticTransformFriendlyModInt() = default;

  static void init() {
    if (roots.empty()) {
      const unsigned mod = Mint::mod();
      assert(mod >= 3 && mod % 2 == 1);
      auto tmp = mod - 1;
      max_base = 0;
      while (tmp % 2 == 0) tmp >>= 1, max_base++;
      Mint root = 2;
      while (root.pow((mod - 1) >> 1) == 1) {
        root += 1;
      }
      assert(root.pow(mod - 1) == 1);

      roots.resize(max_base + 1);
      iroots.resize(max_base + 1);
      rate3.resize(max_base + 1);
      irate3.resize(max_base + 1);

      roots[max_base] = root.pow((mod - 1) >> max_base);
      iroots[max_base] = Mint(1) / roots[max_base];
      for (int i = max_base - 1; i >= 0; i--) {
        roots[i] = roots[i + 1] * roots[i + 1];
        iroots[i] = iroots[i + 1] * iroots[i + 1];
      }
      {
        Mint prod = 1, iprod = 1;
        for (int i = 0; i <= max_base - 3; i++) {
          rate3[i] = roots[i + 3] * prod;
          irate3[i] = iroots[i + 3] * iprod;
          prod *= iroots[i + 3];
          iprod *= roots[i + 3];
        }
      }
    }
  }

  static void ntt(vector<Mint> &a) {
    init();
    const int n = (int)a.size();
    assert((n & (n - 1)) == 0);
    int h = __builtin_ctz(n);
    assert(h <= max_base);
    int len = 0;
    Mint imag = roots[2];
    if (h & 1) {
      int p = 1 << (h - 1);
      Mint rot = 1;
      for (int i = 0; i < p; i++) {
        auto r = a[i + p];
        a[i + p] = a[i] - r;
        a[i] += r;
      }
      len++;
    }
    for (; len + 1 < h; len += 2) {
      int p = 1 << (h - len - 2);
      {  // s = 0
        for (int i = 0; i < p; i++) {
          auto a0 = a[i];
          auto a1 = a[i + p];
          auto a2 = a[i + 2 * p];
          auto a3 = a[i + 3 * p];
          auto a1na3imag = (a1 - a3) * imag;
          auto a0a2 = a0 + a2;
          auto a1a3 = a1 + a3;
          auto a0na2 = a0 - a2;
          a[i] = a0a2 + a1a3;
          a[i + 1 * p] = a0a2 - a1a3;
          a[i + 2 * p] = a0na2 + a1na3imag;
          a[i + 3 * p] = a0na2 - a1na3imag;
        }
      }
      Mint rot = rate3[0];
      for (int s = 1; s < (1 << len); s++) {
        int offset = s << (h - len);
        Mint rot2 = rot * rot;
        Mint rot3 = rot2 * rot;
        for (int i = 0; i < p; i++) {
          auto a0 = a[i + offset];
          auto a1 = a[i + offset + p] * rot;
          auto a2 = a[i + offset + 2 * p] * rot2;
          auto a3 = a[i + offset + 3 * p] * rot3;
          auto a1na3imag = (a1 - a3) * imag;
          auto a0a2 = a0 + a2;
          auto a1a3 = a1 + a3;
          auto a0na2 = a0 - a2;
          a[i + offset] = a0a2 + a1a3;
          a[i + offset + 1 * p] = a0a2 - a1a3;
          a[i + offset + 2 * p] = a0na2 + a1na3imag;
          a[i + offset + 3 * p] = a0na2 - a1na3imag;
        }
        rot *= rate3[__builtin_ctz(~s)];
      }
    }
  }

  static void intt(vector<Mint> &a, bool f = true) {
    init();
    const int n = (int)a.size();
    assert((n & (n - 1)) == 0);
    int h = __builtin_ctz(n);
    assert(h <= max_base);
    int len = h;
    Mint iimag = iroots[2];
    for (; len > 1; len -= 2) {
      int p = 1 << (h - len);
      {  // s = 0
        for (int i = 0; i < p; i++) {
          auto a0 = a[i];
          auto a1 = a[i + 1 * p];
          auto a2 = a[i + 2 * p];
          auto a3 = a[i + 3 * p];
          auto a2na3iimag = (a2 - a3) * iimag;
          auto a0na1 = a0 - a1;
          auto a0a1 = a0 + a1;
          auto a2a3 = a2 + a3;
          a[i] = a0a1 + a2a3;
          a[i + 1 * p] = (a0na1 + a2na3iimag);
          a[i + 2 * p] = (a0a1 - a2a3);
          a[i + 3 * p] = (a0na1 - a2na3iimag);
        }
      }
      Mint irot = irate3[0];
      for (int s = 1; s < (1 << (len - 2)); s++) {
        int offset = s << (h - len + 2);
        Mint irot2 = irot * irot;
        Mint irot3 = irot2 * irot;
        for (int i = 0; i < p; i++) {
          auto a0 = a[i + offset];
          auto a1 = a[i + offset + 1 * p];
          auto a2 = a[i + offset + 2 * p];
          auto a3 = a[i + offset + 3 * p];
          auto a2na3iimag = (a2 - a3) * iimag;
          auto a0na1 = a0 - a1;
          auto a0a1 = a0 + a1;
          auto a2a3 = a2 + a3;
          a[i + offset] = a0a1 + a2a3;
          a[i + offset + 1 * p] = (a0na1 + a2na3iimag) * irot;
          a[i + offset + 2 * p] = (a0a1 - a2a3) * irot2;
          a[i + offset + 3 * p] = (a0na1 - a2na3iimag) * irot3;
        }
        irot *= irate3[__builtin_ctz(~s)];
      }
    }
    if (len >= 1) {
      int p = 1 << (h - 1);
      for (int i = 0; i < p; i++) {
        auto ajp = a[i] - a[i + p];
        a[i] += a[i + p];
        a[i + p] = ajp;
      }
    }
    if (f) {
      Mint inv_sz = Mint(1) / n;
      for (int i = 0; i < n; i++) a[i] *= inv_sz;
    }
  }

  static vector<Mint> multiply(vector<Mint> a, vector<Mint> b) {
    int need = a.size() + b.size() - 1;
    int nbase = 1;
    while ((1 << nbase) < need) nbase++;
    int sz = 1 << nbase;
    a.resize(sz, 0);
    b.resize(sz, 0);
    ntt(a);
    ntt(b);
    Mint inv_sz = Mint(1) / sz;
    for (int i = 0; i < sz; i++) a[i] *= b[i] * inv_sz;
    intt(a, false);
    a.resize(need);
    return a;
  }
};

template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::roots =
    vector<Mint>();
template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::iroots =
    vector<Mint>();
template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::rate3 =
    vector<Mint>();
template <typename Mint>
vector<Mint> NumberTheoreticTransformFriendlyModInt<Mint>::irate3 =
    vector<Mint>();
template <typename Mint>
int NumberTheoreticTransformFriendlyModInt<Mint>::max_base = 0;

Back to top page