普通的FNTT,输入两个多项式A[]和B[],A[i]表示多项式A的x^i的系数。
然后直接Convoluton,默认是使用A[]存放输出的结果,B[]可以视情况回收。
namespace FNTT {
const int MOD = 998244353;
const int G = 3;
inline int add(const int &x, const int &y) {
int r = x + y;
if(r >= MOD)
r -= MOD;
return r;
}
inline int sub(const int &x, const int &y) {
int r = x - y;
if(r < 0)
r += MOD;
return r;
}
inline int mul(const int &x, const int &y) {
ll r = (ll)x * y;
if(r >= MOD)
r %= MOD;
return (int)r;
}
int pow(ll x, int n) {
ll res = 1;
while(n) {
if(n & 1)
res = mul(res, x);
x = mul(x, x);
n >>= 1;
}
return res;
}
void FNTT(vector<int>& a, int n, int op) {
for(int i = 1, j = n >> 1, k; i < n - 1; ++i, j += k) {
if(i < j)
swap(a[i], a[j]);
for(k = n >> 1; k <= j; j -= k, k >>= 1);
}
for(int l = 1; (1 << l) <= n; ++l) {
for(int i = 0, Gl = pow(G, (MOD - 1) / (1 << l)); i < n; i += (1 << l)) {
for(int j = i, w = 1; j < i + (1 << (l - 1)); ++j, w = mul(w, Gl)) {
int u = a[j], t = mul(a[j + (1 << (l - 1))], w);
a[j] = add(u, t), a[j + (1 << (l - 1))] = sub(u, t);
}
}
}
if(op == -1) {
reverse(next(a.begin()), a.end());
for(int i = 0, inv = pow(n, MOD - 2); i < n; ++i)
a[i] = mul(a[i], inv);
}
}
void Convolution(vector<int> &A, vector<int> &B) {
int Asize = A.size(), Bsize = B.size(), n;
for(n = 1; n < (Asize + Bsize - 1); n <<= 1);
A.resize(n), B.resize(n);
FNTT(A, n, 1), FNTT(B, n, 1);
for(int i = 0; i < n; ++i)
A[i] = mul(A[i], B[i]);
FNTT(A, n, -1);
A.resize(Asize + Bsize - 1);
}
}
也可以把Convolution的结果放到一个新的vi里。
vi Convolution(vi &A, vi &B) {
int Asize = A.size(), Bsize = B.size(), n;
for(n = 1; n < (Asize + Bsize - 1); n <<= 1);
A.resize(n), B.resize(n);
FNTT(A, n, 1), FNTT(B, n, 1);
vi C(n);
for(int i = 0; i < n; ++i)
C[i] = mul(A[i], B[i]);
FNTT(C, n, -1);
C.resize(Asize + Bsize - 1);
return move(C);
}
可以把要合并的多项式全部放在一个vec数组里,然后分治进行合并,最后的结果放在vec[1]中。(用于高达n个多项式乘起来,控制长为n的多项式至多作logn次FNTT,总复杂度nlog^2n)
void DCFNTT(int l, int r) {
if(l == r)
return;
int m = (l + r) >> 1;
DCFNTT(l, m), DCFNTT(m + 1, r);
Convolution(vec[l], vec[m + 1]);
}
还可以做的优化:适应更多模数(自带找原根)、FNTT截断:在规模较小的两个多项式相乘使用暴力。