(自用)gpt 大人写的 ntt 基础内容板子

使用模型:6.1 sol high

ntt_poly.hpp:

#pragma once
#include <bits/stdc++.h>

// C++17, modulus 998244353, primitive root 3. No CPU-specific instructions.
// Coefficients are stored from low to high degree.
// raw(x) / x.val() returns the ordinary residue; internal storage is Montgomery.
template <unsigned Mod>
class modint {
  static_assert(Mod > 2 && Mod < (1u << 30) && (Mod & 1));
  using u32 = uint32_t;
  using u64 = uint64_t;
  u32 v_ = 0; // Lazy Montgomery residue in [0, 2*Mod).
  static constexpr u32 neg_inv = [] {
    u32 r = Mod;
    for (int i = 0; i < 5; ++i) r *= 2u - Mod * r;
    return 0u - r;
  }();
  static constexpr u64 R = (u64(1) << 32) % Mod;
  static constexpr u32 R2 = R * R % Mod;
  static u32 reduce(u64 x) {
    return (x + u64(u32(x) * neg_inv) * Mod) >> 32;
  }

public:
  static constexpr int mod = Mod;
  modint() = default;
  template <class T, std::enable_if_t<std::is_integral_v<T>, int> = 0>
  modint(T x) {
    u32 y;
    if constexpr (std::is_signed_v<T>) {
      int64_t t = int64_t(x) % int64_t(Mod);
      if (t < 0) t += Mod;
      y = u32(t);
    } else {
      y = u64(x) % Mod;
    }
    v_ = reduce(u64(y) * R2);
  }
  modint(const std::string& s) {
    assert(!s.empty());
    size_t i = (s[0] == '-' || s[0] == '+');
    assert(i < s.size());
    for (; i < s.size(); ++i) {
      assert('0' <= s[i] && s[i] <= '9');
      *this = *this * 10 + (s[i] - '0');
    }
    if (s[0] == '-') *this = -*this;
  }
  int val() const {
    u32 x = reduce(v_);
    return x >= Mod ? x - Mod : x;
  }
  friend int raw(const modint& x) { return x.val(); }
  modint operator+() const { return *this; }
  modint operator-() const { return modint() - *this; }
  modint& operator+=(const modint& b) {
    v_ += b.v_;
    if (v_ >= 2 * Mod) v_ -= 2 * Mod;
    return *this;
  }
  modint& operator-=(const modint& b) {
    v_ += 2 * Mod - b.v_;
    if (v_ >= 2 * Mod) v_ -= 2 * Mod;
    return *this;
  }
  modint& operator*=(const modint& b) {
    v_ = reduce(u64(v_) * b.v_);
    return *this;
  }
  template <class T> friend modint qpow(modint a, T e) {
    assert(e >= 0);
    modint r = 1;
    for (; e; e >>= 1, a *= a) if (e & 1) r *= a;
    return r;
  }
  // Requires a prime modulus and a nonzero argument.
  modint inv() const { assert(val() != 0); return qpow(*this, Mod - 2); }
  modint& operator/=(const modint& b) { return *this *= b.inv(); }
  friend modint operator+(modint a, const modint& b) { return a += b; }
  friend modint operator-(modint a, const modint& b) { return a -= b; }
  friend modint operator*(modint a, const modint& b) { return a *= b; }
  friend modint operator/(modint a, const modint& b) { return a /= b; }
  friend bool operator==(const modint& a, const modint& b) {
    u32 x = a.v_ >= Mod ? a.v_ - Mod : a.v_;
    u32 y = b.v_ >= Mod ? b.v_ - Mod : b.v_;
    return x == y;
  }
  friend bool operator!=(const modint& a, const modint& b) { return !(a == b); }
  friend std::istream& operator>>(std::istream& in, modint& a) {
    std::string s;
    if (in >> s) a = modint(s);
    return in;
  }
  friend std::ostream& operator<<(std::ostream& out, const modint& a) {
    return out << a.val();
  }
};
using mint = modint<998244353>;

inline int glim(int x) {
  assert(0 <= x && x <= (1 << 30));
  return x <= 1 ? 1 : 1u << (32 - __builtin_clz(unsigned(x - 1)));
}

namespace poly_detail {
inline constexpr int max_ntt = 1 << 23;
inline constexpr int naive_limit = 32;
inline constexpr int fps_base = 32; // Must be a power of two.

// DIF forward and DIT inverse share bit-reversed frequency order.
// Do not mix this forward transform with a natural-order inverse NTT.
inline void ntt(mint* a, int n, bool inverse) {
  assert(n >= 1 && n <= max_ntt && !(n & (n - 1)));
  static std::vector<mint> roots{0, 1}, iroots{0, 1};
  for (int k = int(roots.size()); k < n; k <<= 1) {
    roots.resize(k << 1);
    iroots.resize(k << 1);
    mint w = qpow(mint(3), (mint::mod - 1) / (k << 1));
    mint iw = w.inv();
    roots[k] = iroots[k] = 1;
    for (int j = 1; j < k; ++j) {
      roots[k + j] = roots[k + j - 1] * w;
      iroots[k + j] = iroots[k + j - 1] * iw;
    }
  }
  if (!inverse) {
    // Fuse two radix-2 levels: intermediates stay in registers.
    int k = n >> 1;
    for (; k >= 2; k >>= 2) {
      int q = k >> 1;
      for (int i = 0; i < n; i += q << 2)
        for (int j = 0; j < q; ++j) {
          mint a0 = a[i + j], a1 = a[i + j + q];
          mint a2 = a[i + j + 2 * q], a3 = a[i + j + 3 * q];
          mint x0 = a0 + a2, x1 = a1 + a3;
          mint x2 = (a0 - a2) * roots[k + j];
          mint x3 = (a1 - a3) * roots[k + q + j];
          mint w = roots[q + j];
          a[i + j] = x0 + x1;
          a[i + j + q] = (x0 - x1) * w;
          a[i + j + 2 * q] = x2 + x3;
          a[i + j + 3 * q] = (x2 - x3) * w;
        }
    }
    if (k) for (int i = 0; i < n; i += 2) {
      mint x = a[i], y = a[i + 1];
      a[i] = x + y; a[i + 1] = x - y;
    }
  } else {
    int q = 1;
    for (; (q << 2) <= n; q <<= 2)
      for (int i = 0; i < n; i += q << 2)
        for (int j = 0; j < q; ++j) {
          mint w = iroots[q + j];
          mint a0 = a[i + j], a1 = a[i + j + q] * w;
          mint a2 = a[i + j + 2 * q], a3 = a[i + j + 3 * q] * w;
          mint x0 = a0 + a1, x1 = a0 - a1;
          mint x2 = (a2 + a3) * iroots[2 * q + j];
          mint x3 = (a2 - a3) * iroots[3 * q + j];
          a[i + j] = x0 + x2;
          a[i + j + q] = x1 + x3;
          a[i + j + 2 * q] = x0 - x2;
          a[i + j + 3 * q] = x1 - x3;
        }
    if (q < n) for (int j = 0; j < q; ++j) {
      mint x = a[j], y = a[j + q] * iroots[q + j];
      a[j] = x + y; a[j + q] = x - y;
    }
    static const std::array<mint, 24> inv_pow2 = [] {
      std::array<mint, 24> v{};
      v[0] = 1;
      for (int i = 1; i < 24; ++i) v[i] = v[i - 1] * mint((mint::mod + 1) / 2);
      return v;
    }();
    mint invn = inv_pow2[__builtin_ctz(unsigned(n))];
    for (int i = 0; i < n; ++i) a[i] *= invn;
  }
}
inline const std::vector<mint>& integer_inverses(int n) {
  assert(0 <= n && n < mint::mod);
  static std::vector<mint> inv{0, 1};
  for (int i = int(inv.size()); i <= n; ++i)
    inv.push_back(-mint(mint::mod / i) * inv[mint::mod % i]);
  return inv;
}
} // namespace poly_detail

struct poly : std::vector<mint> {
  using std::vector<mint>::vector;
  poly() = default;
  poly(const std::vector<mint>& a) : std::vector<mint>(a) {}
  poly(std::vector<mint>&& a) : std::vector<mint>(std::move(a)) {}
  // cut() preserves the original template's resize-and-zero-pad semantics.
  poly& cut(int n) { assert(n >= 0); resize(n); return *this; }
  poly pre(int n) const {
    assert(n >= 0);
    poly b(begin(), begin() + std::min(size_t(n), size()));
    b.resize(n);
    return b;
  }
  poly& trim() { while (!empty() && back() == 0) pop_back(); return *this; }
  mint operator()(mint x) const {
    mint y = 0;
    for (int i = int(size()) - 1; i >= 0; --i) y = y * x + (*this)[i];
    return y;
  }
  void ntt(int op) { assert(op == 1 || op == -1); poly_detail::ntt(data(), int(size()), op == -1); }
  poly inv(int n) const;
  poly deriv() const;
  poly integ() const;
  poly ln(int n) const;
  poly exp(int n) const;
  poly pow(long long k, int n) const;
};
inline std::istream& operator>>(std::istream& in, poly& a) {
  for (auto& x : a) in >> x;
  return in;
}
inline std::ostream& operator<<(std::ostream& out, const poly& a) {
  for (size_t i = 0; i < a.size(); ++i) { if (i) out << ' '; out << a[i]; }
  return out;
}
inline poly& operator+=(poly& a, const poly& b) {
  if (a.size() < b.size()) a.resize(b.size());
  for (size_t i = 0; i < b.size(); ++i) a[i] += b[i];
  return a;
}
inline poly& operator-=(poly& a, const poly& b) {
  if (a.size() < b.size()) a.resize(b.size());
  for (size_t i = 0; i < b.size(); ++i) a[i] -= b[i];
  return a;
}
inline poly& operator*=(poly& a, mint c) {
  for (auto& x : a) x *= c;
  return a;
}
inline poly& operator/=(poly& a, mint c) { return a *= c.inv(); }
inline poly& operator<<=(poly& a, int k) {
  assert(k >= 0);
  if (!a.empty()) a.insert(a.begin(), k, 0);
  return a;
}
inline poly& operator>>=(poly& a, int k) {
  assert(k >= 0);
  a.erase(a.begin(), a.begin() + std::min(size_t(k), a.size()));
  return a;
}
inline poly operator+(poly a, const poly& b) { a += b; return a; }
inline poly operator-(poly a, const poly& b) { a -= b; return a; }
inline poly operator-(poly a) { for (auto& x : a) x = -x; return a; }
inline poly operator*(poly a, mint c) { a *= c; return a; }
inline poly operator*(mint c, poly a) { a *= c; return a; }
inline poly operator/(poly a, mint c) { a /= c; return a; }
inline poly operator<<(poly a, int k) { a <<= k; return a; }
inline poly operator>>(poly a, int k) { a >>= k; return a; }

// Returns min(lim, a.size()+b.size()-1) coefficients; does not pad to lim.
// lim == -1 means full convolution. A*A automatically uses one forward NTT.
inline poly mul(const poly& a, const poly& b, int lim = -1) {
  assert(lim >= -1);
  if (a.empty() || b.empty() || lim == 0) return {};
  int na = int(a.size()), nb = int(b.size());
  if (lim >= 0) na = std::min(na, lim), nb = std::min(nb, lim);
  int len = na + nb - 1;
  int need = lim < 0 ? len : std::min(lim, len);
  if (std::min(na, nb) <= poly_detail::naive_limit) {
    poly c(need);
    for (int i = 0; i < na; ++i)
      for (int j = 0; j < nb && i + j < need; ++j) c[i + j] += a[i] * b[j];
    return c;
  }
  assert(len <= poly_detail::max_ntt);
  int n = glim(len);
  poly x(a.begin(), a.begin() + na);
  x.resize(n);
  x.ntt(1);
  if (&a == &b) {
    for (auto& v : x) v *= v;
  } else {
    poly y(b.begin(), b.begin() + nb);
    y.resize(n);
    y.ntt(1);
    for (int i = 0; i < n; ++i) x[i] *= y[i];
  }
  x.ntt(-1);
  x.resize(need);
  return x;
}
inline poly operator*(const poly& a, const poly& b) { return mul(a, b); }
inline poly& operator*=(poly& a, const poly& b) { a = a * b; return a; }

// Formal power series inverse modulo x^n. Requires a[0] != 0.
// Each doubling uses five length-2m transforms, reusing DFT(b).
inline poly getInv(const poly& a, int n) {
  assert(0 <= n && n <= poly_detail::max_ntt);
  if (n == 0) return {};
  assert(!a.empty() && a[0] != 0);
  int base = std::min(n, poly_detail::fps_base);
  poly b(base);
  mint c = a[0].inv();
  b[0] = c;
  for (int i = 1; i < base; ++i) {
    mint s = 0;
    for (int j = 1; j <= i && j < int(a.size()); ++j) s += a[j] * b[i - j];
    b[i] = -s * c;
  }
  while (int(b.size()) < n) {
    int m = int(b.size()), len = m << 1, k = std::min(len, n);
    poly x = a.pre(k), y = b;
    x.resize(len);
    y.resize(len);
    x.ntt(1);
    y.ntt(1);
    for (int i = 0; i < len; ++i) x[i] *= y[i];
    x.ntt(-1);
    // Cyclic wraparound only affects indices below m; discard them.
    std::fill(x.begin(), x.begin() + m, mint(0));
    std::fill(x.begin() + k, x.end(), mint(0));
    x.ntt(1);
    for (int i = 0; i < len; ++i) x[i] *= y[i];
    x.ntt(-1);
    b.resize(k);
    for (int i = m; i < k; ++i) b[i] = -x[i];
  }
  return b;
}
inline poly getDer(const poly& a) {
  if (a.empty()) return {};
  poly b(a.size() - 1);
  for (int i = 1; i < int(a.size()); ++i) b[i - 1] = a[i] * i;
  return b;
}
inline poly getInt(const poly& a) {
  const auto& inv = poly_detail::integer_inverses(int(a.size()));
  poly b(a.size() + 1);
  for (int i = 0; i < int(a.size()); ++i) b[i + 1] = a[i] * inv[i + 1];
  return b;
}
// Formal logarithm modulo x^n, constant term 0. Requires a[0] == 1.
inline poly getLn(const poly& a, int n) {
  assert(0 <= n && n <= poly_detail::max_ntt);
  if (!n) return {};
  assert(!a.empty() && a[0] == 1);
  if (n == 1) return poly(1);
  poly d = getDer(poly(a.begin(), a.begin() + std::min(a.size(), size_t(n))));
  poly b = getInt(mul(d, getInv(a, n - 1), n - 1));
  b.resize(n);
  return b;
}

// Formal exponential modulo x^n. Requires a.empty() || a[0] == 0.
// Maintains b = exp(a) and c = 1/b modulo x^m.
// Each doubling uses length-2m transforms, not length-4m convolutions.
inline poly getExp(const poly& a, int n) {
  assert(0 <= n && n <= poly_detail::max_ntt);
  if (!n) return {};
  assert(a.empty() || a[0] == 0);
  const auto& inv = poly_detail::integer_inverses(n - 1);
  int base = std::min(n, poly_detail::fps_base);
  poly b(base);
  b[0] = 1;
  for (int i = 1; i < base; ++i) {
    mint s = 0;
    for (int j = 1; j <= i && j < int(a.size()); ++j) s += mint(j) * a[j] * b[i - j];
    b[i] = s * inv[i];
  }
  if (base == n) return b;
  poly c = getInv(b, base);
  while (int(b.size()) < n) {
    int m = int(b.size()), len = m << 1, k = std::min(len, n);
    poly B = b, C = c, H(len);
    B.resize(len);
    C.resize(len);
    B.ntt(1);
    C.ntt(1);
    // H is the known low part of (log b)' = a' mod x^(m-1).
    for (int i = 0; i < m - 1 && i + 1 < int(a.size()); ++i) H[i] = a[i + 1] * (i + 1);
    H.ntt(1);
    for (int i = 0; i < len; ++i) H[i] *= B[i];
    H.ntt(-1);
    std::fill(H.begin(), H.begin() + m - 1, mint(0));
    for (int i = m - 1; i < k - 1; ++i) H[i] = -H[i];
    std::fill(H.begin() + k - 1, H.end(), mint(0));
    H.ntt(1);
    for (int i = 0; i < len; ++i) H[i] *= C[i];
    H.ntt(-1);
    // Now H[m-1..k-2] is the high part of (log b)'.
    poly D(len);
    for (int i = m; i < k; ++i)
      D[i] = (i < int(a.size()) ? a[i] : mint(0)) - H[i - 1] * inv[i];
    if (k < n) {
      // c_new = c - c*(b*c - 1 + D) mod x^k.
      // The parenthesized error starts at degree m.
      for (int i = 0; i < len; ++i) H[i] = B[i] * C[i];
      H.ntt(-1);
      std::fill(H.begin(), H.begin() + m, mint(0));
      for (int i = m; i < k; ++i) H[i] += D[i];
      std::fill(H.begin() + k, H.end(), mint(0));
      H.ntt(1);
      for (int i = 0; i < len; ++i) H[i] *= C[i];
      H.ntt(-1);
      c.resize(k);
      for (int i = m; i < k; ++i) c[i] = -H[i];
    }
    D.ntt(1);
    for (int i = 0; i < len; ++i) D[i] *= B[i];
    D.ntt(-1);
    b.resize(k);
    for (int i = m; i < k; ++i) b[i] = D[i];
  }
  return b;
}

// Nonnegative integer power modulo x^n; handles leading zeros and a[0] != 1.
inline poly getPow(const poly& a, long long k, int n) {
  assert(k >= 0 && 0 <= n && n <= poly_detail::max_ntt);
  if (!n) return {};
  if (!k) { poly b(n); b[0] = 1; return b; }
  int v = 0;
  while (v < int(a.size()) && a[v] == 0) ++v;
  if (v == int(a.size()) || v > (n - 1) / k) return poly(n);
  int shift = int(v * k), len = n - shift;
  mint c = a[v];
  poly b(a.begin() + v, a.begin() + v + std::min(int(a.size()) - v, len));
  b /= c;
  b = getExp(getLn(b, len) * mint(k), len) * qpow(c, k);
  b <<= shift;
  b.resize(n);
  return b;
}

namespace poly_detail {
// Ordinary polynomial division. Leading (high-degree) zeros are ignored.
inline std::pair<poly, poly> divide(poly a, poly b, bool need_remainder) {
  a.trim(); b.trim();
  assert(!b.empty());
  if (a.size() < b.size()) return {{}, need_remainder ? std::move(a) : poly{}};
  int na = int(a.size()), nb = int(b.size()), nq = na - nb + 1;
  poly q(nq);
  if (std::min(nq, nb) <= naive_limit) {
    mint iv = b.back().inv();
    for (int i = nq - 1; i >= 0; --i) {
      mint t = q[i] = a[i + nb - 1] * iv;
      for (int j = 0; j < nb; ++j) a[i + j] -= t * b[j];
    }
  } else {
    poly ar(a.rbegin(), a.rbegin() + nq);
    poly br(b.rbegin(), b.rbegin() + std::min(nb, nq));
    q = mul(ar, getInv(br, nq), nq);
    std::reverse(q.begin(), q.end());
    if (need_remainder) a -= b * q;
  }
  if (!need_remainder) return {std::move(q), {}};
  a.resize(nb - 1);
  a.trim();
  return {std::move(q), std::move(a)};
}
} // namespace poly_detail
inline std::pair<poly, poly> divmod(const poly& a, const poly& b) {
  return poly_detail::divide(a, b, true);
}
inline poly operator/(const poly& a, const poly& b) { return poly_detail::divide(a, b, false).first; }
inline poly operator%(const poly& a, const poly& b) { return poly_detail::divide(a, b, true).second; }
inline poly& operator/=(poly& a, const poly& b) { a = a / b; return a; }
inline poly& operator%=(poly& a, const poly& b) { a = a % b; return a; }
inline poly poly::inv(int n) const { return getInv(*this, n); }
inline poly poly::deriv() const { return getDer(*this); }
inline poly poly::integ() const { return getInt(*this); }
inline poly poly::ln(int n) const { return getLn(*this, n); }
inline poly poly::exp(int n) const { return getExp(*this, n); }
inline poly poly::pow(long long k, int n) const { return getPow(*this, k, n); }

ntt_poly_check.cpp:

#include "ntt_poly.hpp"
using namespace std;
static mt19937_64 rng(20261020);
static mint rnd() { return mint(rng() % mint::mod); }
static poly random_poly(int n) { poly a(n); for (auto& x : a) x = rnd(); return a; }
static poly naive_mul(const poly& a, const poly& b, int lim = -1) {
  if (a.empty() || b.empty() || lim == 0) return {};
  int n = int(a.size() + b.size() - 1);
  if (lim >= 0) n = min(n, lim);
  poly c(n);
  for (int i = 0; i < int(a.size()) && i < n; ++i)
    for (int j = 0; j < int(b.size()) && i + j < n; ++j) c[i + j] += a[i] * b[j];
  return c;
}
static poly naive_inv(const poly& a, int n) {
  poly b(n);
  if (!n) return b;
  mint iv = a[0].inv(); b[0] = iv;
  for (int i = 1; i < n; ++i) {
    mint s = 0;
    for (int j = 1; j <= i && j < int(a.size()); ++j) s += a[j] * b[i-j];
    b[i] = -s * iv;
  }
  return b;
}
static poly naive_exp(const poly& a, int n) {
  poly b(n);
  if (!n) return b;
  b[0] = 1;
  for (int i = 1; i < n; ++i) {
    for (int j = 1; j <= i && j < int(a.size()); ++j) b[i] += mint(j) * a[j] * b[i-j];
    b[i] /= mint(i);
  }
  return b;
}
static poly naive_ln(const poly& a, int n) {
  if (!n) return {};
  poly d = getDer(a.pre(n));
  poly b = getInt(naive_mul(d, naive_inv(a, n), n - 1));
  b.resize(n); return b;
}
static pair<poly,poly> naive_div(poly a, poly b) {
  a.trim(); b.trim(); assert(!b.empty());
  if (a.size() < b.size()) return {{}, a};
  int n = int(a.size()), m = int(b.size()); poly q(n-m+1);
  mint iv = b.back().inv();
  for (int i=n-m; i>=0; --i) {
    q[i] = a[i+m-1] * iv;
    for (int j=0; j<m; ++j) a[i+j] -= q[i]*b[j];
  }
  a.resize(m-1); a.trim(); return {q,a};
}
static void check(bool ok, const char* msg) {
  if (!ok) { cerr << "FAILED: " << msg << '\n'; exit(1); }
}
static void tests() {
  const uint64_t p = mint::mod;
  vector<uint64_t> edge{0,1,2,p-2,p-1,p,p+1,UINT64_MAX};
  for (int t=0; t<200000; ++t) {
    uint64_t a = t < int(edge.size()) ? edge[t] : rng();
    uint64_t b = t < int(edge.size()) ? edge[edge.size()-1-t] : rng();
    mint x(a), y(b);
    check(raw(x)==int(a%p), "mint constructor");
    check(raw(x+y)==int((a%p+b%p)%p), "mint add");
    check(raw(x-y)==int((a%p+p-b%p)%p), "mint sub");
    check(raw(x*y)==int((a%p)*(b%p)%p), "mint mul");
    check(x+(-x)==0 && x==mint(raw(x)), "mint lazy equality");
    if (t<5000 && y!=0) check(x/y*y==x, "mint division");
  }
  for (long long a : {LLONG_MIN, -998244354LL, -1LL, 0LL, 1LL, LLONG_MAX}) {
    long long v = a % mint::mod; if (v<0) v+=mint::mod;
    check(raw(mint(a))==v, "signed constructor");
  }
  check(mint(string("-123456789012345678901234567890")) == -mint(string("123456789012345678901234567890")), "string constructor");
  for (int k=0; k<=18; ++k) {
    poly a=random_poly(1<<k), b=a; b.ntt(1); b.ntt(-1);
    check(a==b,"NTT round trip");
  }
  vector<int> lens{0,1,2,3,15,16,17,31,32,33,63,64,65,127,128,129,255,256,257,511,512,513};
  for (int na:lens) for (int nb:lens) {
    poly a=random_poly(na), b=random_poly(nb);
    check(a*b==naive_mul(a,b),"convolution");
    int lim=int(rng()%700);
    check(mul(a,b,lim)==naive_mul(a,b,lim),"truncated convolution");
    check(a*a==naive_mul(a,a),"square");
  }
  for (int t=0; t<300; ++t) {
    int n=t<int(lens.size())?lens[t]:int(rng()%400);
    poly a=random_poly(1+int(rng()%500)); if (a[0]==0) a[0]=1;
    check(getInv(a,n)==naive_inv(a,n),"inverse recurrence");
    a[0]=1;
    check(getLn(a,n)==naive_ln(a,n),"log recurrence");
    a[0]=0;
    check(getExp(a,n)==naive_exp(a,n),"exp recurrence");
    check(getLn(getExp(a,n),n)==a.pre(n),"ln(exp(a))");
  }
  for (int n : lens) {
    check(getExp({},n)==poly{1}.pre(n),"exp zero");
    check(getInv(poly{7},n)==poly{mint(7).inv()}.pre(n),"inverse constant");
  }
  for (int t=0; t<300; ++t) {
    int na=int(rng()%400), nb=1+int(rng()%200);
    poly a=random_poly(na), b=random_poly(nb);
    if (b.back()==0) b.back()=1;
    if (t%3==0) {a.resize(na+3); b.resize(nb+4);}
    auto expected=naive_div(a,b), got=divmod(a,b);
    check(got==expected,"division and remainder");
    check(a/b==expected.first && a%b==expected.second,"division operators");
  }
  for (int t=0; t<150; ++t) {
    int n=1+int(rng()%150), k=int(rng()%10), sz=int(rng()%100);
    poly a=random_poly(sz);
    if (sz && t%2==0) fill(a.begin(),a.begin()+int(rng()%sz),mint(0));
    poly expected{1};
    for (int i=0; i<k; ++i) expected=naive_mul(expected,a,n);
    expected.resize(n);
    check(getPow(a,k,n)==expected,"power");
  }
  check(getPow(poly{0,1},LLONG_MAX,20)==poly(20),"power shift overflow");
  check(getPow(poly{1,1},mint::mod,20)==poly{1}.pre(20),"power characteristic");
  for (int n : {1024,4093,65536,100000}) {
    poly a=random_poly(n); a[0]=0;
    poly b=getExp(a,n);
    check(getLn(b,n)==a,"large exp/log identity");
    poly c=getInv(b,n), product=mul(b,c,n); product.resize(n);
    check(product==poly{1}.pre(n),"large inverse identity");
  }
  for (int n : {1024,4093,65536}) {
    poly b=random_poly(n), q=random_poly(n+1), r=random_poly(n-1);
    if(b.back()==0) b.back()=1;
    if(q.back()==0) q.back()=1;
    poly a=b*q+r;
    auto got=divmod(a,b);
    r.trim();
    check(got.first==q && got.second==r,"large division identity");
  }
  poly a{1,2,3}, b{4,5};
  check(&(a+=b)==&a && &(a-=b)==&a,"compound assignment returns reference");
  check(a==poly({1,2,3}),"compound assignment values");
  poly alias=a; alias+=alias; check(alias==a*mint(2),"self add");
  alias=a; alias-=alias; check(alias==poly(3),"self sub");
  alias=a; alias*=alias; check(alias==a*a,"self multiply");
  alias=poly{2,3,4}; alias*=alias[0]; check(alias==poly({4,6,8}),"scalar alias");
  alias=a; alias/=alias; check(alias==poly{1},"self divide");
  alias=a; alias%=alias; check(alias.empty(),"self remainder");
  check((a<<2)>>2==a && (a>>99).empty(),"shifts");
  check(getDer(getInt(a))==a,"derivative integral");
  cerr << "All correctness checks passed.\n";
}
static void benchmark(int n) {
  poly a=random_poly(n), b=random_poly(n);
  auto measure=[&](const char* name,auto f) {
    vector<double> times; mint checksum=0;
    for (int t=0;t<4;++t) {
      auto start=chrono::steady_clock::now(); poly c=f();
      auto end=chrono::steady_clock::now();
      checksum+=c[c.size()/2];
      if(t) times.push_back(chrono::duration<double,milli>(end-start).count());
    }
    sort(times.begin(),times.end());
    cout<<name<<' '<<fixed<<setprecision(3)<<times[1]<<" ms; checksum "<<checksum<<'\n';
  };
  measure("multiply",[&]{return a*b;});
  measure("square",[&]{return a*a;});
  a[0]=1;
  measure("inverse",[&]{return getInv(a,n);});
  measure("ln",[&]{return getLn(a,n);});
  a[0]=0;
  measure("exp",[&]{return getExp(a,n);});
}
int main(int argc,char** argv) {
  if (argc>1 && string(argv[1])=="--bench") benchmark(argc>2?stoi(argv[2]):100000);
  else tests();
}

常用 NTT / 多项式模板

主文件:ntt_poly.hpp。C++17,固定使用模数 998244353、原根 3。
可以直接复制到提交代码顶部,也可以使用 #include "ntt_poly.hpp"。
系数按低次到高次存储,例如 poly{1, 2, 3} 表示 1 + 2x + 3x^2。

常用接口

写法 功能 / 约定
mint x; raw(x); x.val(); qpow(x, k) 模整数、读取普通整数值、非负整数快速幂
a + b, a - b, a * b 多项式加减乘;同一个对象的 a * a 自动使用平方优化
a * c, c * a, a / c c 为标量;多项式除以标量时只计算一次逆元
a / b, a % b 普通多项式的商、余数,要求 b 非零
auto [q, r] = divmod(a, b) 同时求商和余数,避免分别调用 /、% 的重复计算
getInv(a, n) / a.inv(n) 形式幂级数逆元模 x^n,要求 a[0] != 0
getDer(a) / a.deriv() 求导
getInt(a) / a.integ() 积分,积分常数为 0
getLn(a, n) / a.ln(n) 形式对数模 x^n,要求 a[0] == 1
getExp(a, n) / a.exp(n) 形式指数模 x^n,要求 a 为空或 a[0] == 0
getPow(a, k, n) / a.pow(k, n) 非负整数幂模 x^n,支持前导零和任意常数项;约定零次幂为 1
mul(a, b, n) 只保留乘积前 n 项;实际返回长度为 min(n, a.size()+b.size()-1),不自动补至 n
a.cut(n) 原地截取或补零到恰好 n 项,返回 poly&,与旧模板语义一致
a.pre(n) 返回前 n 项的副本,不足补零;不修改 a
a.trim() 原地删除高次端的零,零多项式变为空数组
a << k, a >> k 乘 x^k / 去掉前 k 项,要求 k >= 0
a(x) Horner 法计算多项式在 x 处的值
a.ntt(1), a.ntt(-1) 正变换 / 逆变换,长度必须为正的二次幂

相应的 += -= *= /= %= <<= >>= 均返回引用。poly 继续继承 std::vector<mint>,支持原来的下标、迭代器、resize、流输入输出等用法。

getInv/getLn/getExp/getPow 中的 n 是保留的项数,返回值恰好有 n 项;n == 0 时返回空数组。普通加减法保持数组长度,不自动 trim()。普通除法忽略输入高次端的零,商和余数按有效次数返回,零余数为空数组。

#include "ntt_poly.hpp"
int main() {
  poly a{1, 2, 3}, b{4, 5};
  poly product = a * b;
  
  auto [quotient, remainder] = divmod(product, b);
  assert(quotient == a && remainder.empty());
  
  poly inverse = getInv(a, 8);
  assert(mul(a, inverse, 8) == poly{1}.pre(8));
  
  poly f{0, 1, 2};
  poly e = f.exp(8);
  assert(e.ln(8) == f.pre(8));
  
  std::cout << product << '\n';
  printf("%d\n", raw(product[0]));
}

实现与限制

  • Montgomery 模乘,内部允许 [0, 2*mod) 的冗余表示。内部存储不暴露为 .v;读取普通整数用 raw(x) 或 x.val()。不要把 mint 对象直接传给 printf("%d", ...)。
  • DIF 正变换与 DIT 逆变换配对,并合并相邻两层蝶形,减少中间数组读写。正变换结果按位逆序排列;逐点乘法不受影响,但不能与另一份采用自然顺序的 NTT 混用。
  • 单位根与整数逆元按需缓存;没有 CPU 特定指令,也没有在头文件中强加编译优化 pragma。
  • 普通乘法不经过 std::function 和逐点临时容器。平方少做一次正变换。朴素乘法阈值为 32,可修改 poly_detail::naive_limit;FPS 初始块大小由独立的 fps_base 控制,必须是二次幂。
  • 求逆倍增使用 5 次长度 2m 的变换并复用逆多项式的频域值,避免旧实现使用长度 4m 的变换。
  • exp 同时维护指数及其逆元,利用已知低位推算新增高位;最后一轮不更新已经不需要的逆元。各轮使用长度 2m 的变换。
  • 乘法、求逆、ln、exp 都是 O(n log n) 级别;普通多项式除法通过翻转和求逆实现,小规模时使用长除法。
  • 最大 NTT 长度是 2^23。这不等于任意运算都支持 2^23 项输入:普通卷积要满足补齐后的总乘积长度不超过该上限;ln 等调用普通卷积的接口也受中间卷积长度限制。检查由 assert 完成。

本地验证与计时

验证程序是 ntt_poly_check.cpp,与头文件放在同一目录后运行:

g++ -std=c++17 -O2 -Wall -Wextra ntt_poly_check.cpp -o check
./check
./check --bench 100000

验证包括模整数与普通整数取模比较、NTT 正逆变换、乘法/平方与朴素卷积比较、求逆/ln/exp 与朴素递推比较、普通多项式长除法对拍、整数幂与反复乘法比较,以及十万项规模的 exp/ln 和逆元恒等式、较大规模的商余数构造检查。AddressSanitizer 和 UndefinedBehaviorSanitizer 检查通过;当前环境不支持 LeakSanitizer,运行时关闭了泄漏扫描。

posted @ 2026-10-20 11:03  rizynvu  阅读(11)  评论(0)    收藏  举报