math/Poly.h

• View this file on GitHub
• Last update: 2025-11-28 00:00:09+07:00

Depends on

• math/ModInt.h

Verified with

• tests/Bell_Number.test.cpp
• tests/Bernoulli_Number.test.cpp
• tests/Convolution.test.cpp
• tests/Exp_of_FPS.test.cpp
• tests/Inv_of_FPS.test.cpp
• tests/Partition_Function.test.cpp
• tests/Pow_of_FPS.test.cpp
• tests/Product_of_Polynomial_Sequence.test.cpp
• tests/Stirling_Number_1st.test.cpp
• tests/Stirling_Number_1st_fixed_K.test.cpp
• tests/Stirling_Number_2nd.test.cpp
• tests/Stirling_Number_2nd_fixed_K.test.cpp
• tests/p_Subset_Sum.test.cpp

Code

#include "ModInt.h"

using Fp = modint<998244353>;
namespace ntt {
const Fp G = 3;
void ntt(vector<Fp>& a, bool inv) {
  int n = sz(a);
  for (int i = 1, j = 0; i < n; i++) {
    int bit = n >> 1;
    for (; j & bit; bit >>= 1) j ^= bit;
    j ^= bit;
    if (i < j) swap(a[i], a[j]);
  }
  for (int len = 1; len < n; len <<= 1) {
    Fp wlen = G.pow((Fp::modulo - 1) / (2 * len));
    if (inv) wlen = wlen.inv();
    for (int i = 0; i < n; i += 2 * len) {
      Fp w = 1;
      for (int j = 0; j < len; j++) {
        Fp u = a[i + j], v = a[i + j + len] * w;
        a[i + j] = u + v, a[i + j + len] = u - v;
        w *= wlen;
      }
    }
  }
  if (inv) {
    Fp n_inv = Fp(n).inv();
    for (auto& x : a) x *= n_inv;
  }
}
vector<Fp> conv(vector<Fp> a, vector<Fp> b) {
  if (a.empty() || b.empty()) return {};
  int s = sz(a) + sz(b) - 1, n = 1;
  while (n < s) n <<= 1;
  a.resize(n), b.resize(n);
  ntt(a, 0), ntt(b, 0);
  for (int i = 0; i < n; i++) a[i] *= b[i];
  ntt(a, 1), a.resize(s);
  return a;
}
}  // namespace ntt

struct Poly : vector<Fp> {
  using vector::vector;
  Poly(const vector<Fp>& v) : vector(v) {}

  Poly cut(int n) const {
    Poly res = *this;
    res.resize(n);
    return res;
  }

  Poly operator+(const Poly& r) const {
    Poly res = *this;
    res.resize(max(sz(*this), sz(r)));
    for (int i = 0; i < sz(r); ++i) res[i] += r[i];
    return res;
  }
  Poly operator-(const Poly& r) const {
    Poly res = *this;
    res.resize(max(sz(*this), sz(r)));
    for (int i = 0; i < sz(r); ++i) res[i] -= r[i];
    return res;
  }
  Poly operator*(const Poly& r) const { return ntt::conv(*this, r); }
  Poly operator*(Fp v) const {
    Poly res = *this;
    for (auto& x : res) x *= v;
    return res;
  }
  Poly& operator+=(const Poly& r) { return *this = *this + r; }
  Poly& operator-=(const Poly& r) { return *this = *this - r; }
  Poly& operator*=(const Poly& r) { return *this = *this * r; }

  Poly deriv() const {
    if (empty()) return {};
    Poly res(sz(*this) - 1);
    for (int i = 1; i < sz(*this); ++i) res[i - 1] = data()[i] * i;
    return res;
  }
  Poly integ() const {
    Poly res(sz(*this) + 1);
    for (int i = 0; i < sz(*this); ++i) res[i + 1] = data()[i] * Fp(i + 1).inv();
    return res;
  }
  Poly inv(int n) const {
    Poly b = {data()[0].inv()};
    for (int k = 1; k < n; k <<= 1) {
      Poly a = cut(2 * k), prod = b * b * a;
      b.resize(2 * k);
      for (int i = 0; i < 2 * k; ++i) {
        b[i] = b[i] * 2 - (i < sz(prod) ? prod[i] : Fp(0));
      }
    }
    return b.cut(n);
  }

  Poly log(int n) const { return (deriv() * inv(n)).integ().cut(n); }

  Poly exp(int n) const {
    Poly b = {1};
    for (int k = 1; k < n; k <<= 1) {
      Poly ln_b = b.log(2 * k), a = cut(2 * k), diff = a - ln_b;
      diff[0] += 1, b = (b * diff).cut(2 * k);
    }
    return b.cut(n);
  }

  Poly pow(i64 k, int n) const {
    if (n == 0) return {};
    if (k == 0) {
      Poly res = {1};
      res.resize(n);
      return res;
    }
    int i = 0;
    while (i < sz(*this) && data()[i].x == 0) i++;
    if (i == sz(*this) || (i > 0 && k >= n / i + 2)) {
      Poly res;
      res.resize(n);
      return res;
    }
    i64 shift = (i64) i * k;
    if (shift >= n) {
      Poly res;
      res.resize(n);
      return res;
    }
    Poly a = {begin() + i, end()};
    int limit = n - shift;
    a.resize(limit);
    Fp lead = a[0];
    Fp inv_lead = lead.inv();
    a = a * inv_lead;
    a = (a.log(limit) * Fp(k)).exp(limit);
    a = a * lead.pow(k);
    Poly res(shift, 0);
    res.insert(res.end(), a.begin(), a.end());
    res.resize(n);
    return res;
  }
  friend ostream& operator<<(ostream& os, const Poly& p) {
    for (auto x : p) os << x << " ";
    return os;
  }
};

#line 2 "math/ModInt.h"

template <int mod>
struct modint {
  using M = modint;
  static_assert(mod > 0 && mod <= 2147483647);
  static constexpr int modulo = mod;
  static constexpr u32 r1 = []() {
    u32 r1 = mod;
    for (int i = 0; i < 5; ++i) r1 *= 2 - mod * r1;
    return -r1;
  }();
  static constexpr u32 r2 = -u64(mod) % mod;
  static u32 reduce(u64 x) {
    u32 y = u32(x) * r1, r = (x + u64(y) * mod) >> 32;
    return r >= mod ? r - mod : r;
  }
  u32 x;
  modint() : x(0) {}
  modint(i64 x) : x(reduce(u64(x % mod + mod) * r2)) {}
  M& operator+=(const M& a) {
    if ((x += a.x) >= mod) x -= mod;
    return *this;
  }
  M& operator-=(const M& a) {
    if ((x += mod - a.x) >= mod) x -= mod;
    return *this;
  }
  M& operator*=(const M& a) {
    x = reduce(u64(x) * a.x);
    return *this;
  }
  M& operator/=(const M& a) { return *this *= a.inv(); }
  M operator-() const { return M(0) - *this; }
  M operator+(const M& a) const { return M(*this) += a; }
  M operator-(const M& a) const { return M(*this) -= a; }
  M operator*(const M& a) const { return M(*this) *= a; }
  M operator/(const M& a) const { return M(*this) /= a; }
  bool operator==(const M& a) const { return x == a.x; }
  bool operator!=(const M& a) const { return x != a.x; }
  M pow(u64 k) const {
    M res(1), b = *this;
    while (k) {
      if (k & 1) res *= b;
      b *= b, k >>= 1;
    }
    return res;
  }
  M inv() const { return pow(mod - 2); }
  friend ostream& operator<<(ostream& os, const M& a) {
    return os << reduce(a.x);
  }
  friend istream& operator>>(istream& is, M& a) {
    i64 v;
    is >> v;
    a = M(v);
    return is;
  }
};

u64 modmul(u64 x, u64 y, u64 m) { return u128(x) * y % m; }
u64 modpow(u64 x, u64 k, u64 m) {
  u64 res = 1;
  while (k) {
    if (k & 1) res = modmul(res, x, m);
    x = modmul(x, x, m);
    k >>= 1;
  }
  return res;
}
#line 2 "math/Poly.h"

using Fp = modint<998244353>;
namespace ntt {
const Fp G = 3;
void ntt(vector<Fp>& a, bool inv) {
  int n = sz(a);
  for (int i = 1, j = 0; i < n; i++) {
    int bit = n >> 1;
    for (; j & bit; bit >>= 1) j ^= bit;
    j ^= bit;
    if (i < j) swap(a[i], a[j]);
  }
  for (int len = 1; len < n; len <<= 1) {
    Fp wlen = G.pow((Fp::modulo - 1) / (2 * len));
    if (inv) wlen = wlen.inv();
    for (int i = 0; i < n; i += 2 * len) {
      Fp w = 1;
      for (int j = 0; j < len; j++) {
        Fp u = a[i + j], v = a[i + j + len] * w;
        a[i + j] = u + v, a[i + j + len] = u - v;
        w *= wlen;
      }
    }
  }
  if (inv) {
    Fp n_inv = Fp(n).inv();
    for (auto& x : a) x *= n_inv;
  }
}
vector<Fp> conv(vector<Fp> a, vector<Fp> b) {
  if (a.empty() || b.empty()) return {};
  int s = sz(a) + sz(b) - 1, n = 1;
  while (n < s) n <<= 1;
  a.resize(n), b.resize(n);
  ntt(a, 0), ntt(b, 0);
  for (int i = 0; i < n; i++) a[i] *= b[i];
  ntt(a, 1), a.resize(s);
  return a;
}
}  // namespace ntt

struct Poly : vector<Fp> {
  using vector::vector;
  Poly(const vector<Fp>& v) : vector(v) {}

  Poly cut(int n) const {
    Poly res = *this;
    res.resize(n);
    return res;
  }

  Poly operator+(const Poly& r) const {
    Poly res = *this;
    res.resize(max(sz(*this), sz(r)));
    for (int i = 0; i < sz(r); ++i) res[i] += r[i];
    return res;
  }
  Poly operator-(const Poly& r) const {
    Poly res = *this;
    res.resize(max(sz(*this), sz(r)));
    for (int i = 0; i < sz(r); ++i) res[i] -= r[i];
    return res;
  }
  Poly operator*(const Poly& r) const { return ntt::conv(*this, r); }
  Poly operator*(Fp v) const {
    Poly res = *this;
    for (auto& x : res) x *= v;
    return res;
  }
  Poly& operator+=(const Poly& r) { return *this = *this + r; }
  Poly& operator-=(const Poly& r) { return *this = *this - r; }
  Poly& operator*=(const Poly& r) { return *this = *this * r; }

  Poly deriv() const {
    if (empty()) return {};
    Poly res(sz(*this) - 1);
    for (int i = 1; i < sz(*this); ++i) res[i - 1] = data()[i] * i;
    return res;
  }
  Poly integ() const {
    Poly res(sz(*this) + 1);
    for (int i = 0; i < sz(*this); ++i) res[i + 1] = data()[i] * Fp(i + 1).inv();
    return res;
  }
  Poly inv(int n) const {
    Poly b = {data()[0].inv()};
    for (int k = 1; k < n; k <<= 1) {
      Poly a = cut(2 * k), prod = b * b * a;
      b.resize(2 * k);
      for (int i = 0; i < 2 * k; ++i) {
        b[i] = b[i] * 2 - (i < sz(prod) ? prod[i] : Fp(0));
      }
    }
    return b.cut(n);
  }

  Poly log(int n) const { return (deriv() * inv(n)).integ().cut(n); }

  Poly exp(int n) const {
    Poly b = {1};
    for (int k = 1; k < n; k <<= 1) {
      Poly ln_b = b.log(2 * k), a = cut(2 * k), diff = a - ln_b;
      diff[0] += 1, b = (b * diff).cut(2 * k);
    }
    return b.cut(n);
  }

  Poly pow(i64 k, int n) const {
    if (n == 0) return {};
    if (k == 0) {
      Poly res = {1};
      res.resize(n);
      return res;
    }
    int i = 0;
    while (i < sz(*this) && data()[i].x == 0) i++;
    if (i == sz(*this) || (i > 0 && k >= n / i + 2)) {
      Poly res;
      res.resize(n);
      return res;
    }
    i64 shift = (i64) i * k;
    if (shift >= n) {
      Poly res;
      res.resize(n);
      return res;
    }
    Poly a = {begin() + i, end()};
    int limit = n - shift;
    a.resize(limit);
    Fp lead = a[0];
    Fp inv_lead = lead.inv();
    a = a * inv_lead;
    a = (a.log(limit) * Fp(k)).exp(limit);
    a = a * lead.pow(k);
    Poly res(shift, 0);
    res.insert(res.end(), a.begin(), a.end());
    res.resize(n);
    return res;
  }
  friend ostream& operator<<(ostream& os, const Poly& p) {
    for (auto x : p) os << x << " ";
    return os;
  }
};

Back to top page