(自用)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,运行时关闭了泄漏扫描。

浙公网安备 33010602011771号