快速数论变换(NTT)

发布时间 2023-09-16 18:31:48作者: Mcggvc

在系数均为整数的时候,可以用NTT代替FFT,这样不会出现精度问题。

#include <bits/stdc++.h>
using namespace std;
typedef long long lld;
const int N = 20000005;
const lld g = 3, mod = 998244353;
lld r[N];
lld powe(lld a, lld b) {
    lld base = 1;
    while(b) {
        if(b & 1) base = base * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return base;
}
void init(int M) 
{
    int l = log2(M); r[0] = 0; 
    for(int i = 1; i < M; i++) {
        r[i] = (r[i >> 1] >> 1 | (i & 1) << (l - 1));
    }
}
void NTT(vector <lld> &a, int M, int type) {
	for(int i = 0; i < M; i++) {
        if(i<r[i]) swap(a[i], a[r[i]]);
    }
    lld Wn, w;
    for(int mid = 1; mid < M; mid <<= 1) {
        Wn = powe(g, (mod - 1) / (mid << 1));
        if(type == -1) Wn = powe(Wn, mod - 2);
        int size = mid << 1;
        for(int i = 0; i < M; i += size) {
            int j = i + mid; w = 1;
            for(int k = i; k < j; k++) {
                lld x = a[k], y = w * a[k + mid] % mod;
                a[k] = (x + y) % mod;
                a[k + mid] = (x - y + mod) % mod;
                w = w * Wn % mod;
            }
        }
    }
}
vector<lld> conv(vector<lld> a, vector<lld> b) {
	int siz = a.size() + b.size();
	int M = 1;
	while(M < siz) M <<= 1;
	a.resize(M); b.resize(M); 
    init(M); NTT(a, M, 1); NTT(b, M, 1);
    for(int i = 0; i <= M; i++) {
        a[i] = a[i] * b[i] % mod;
    }
    NTT(a, M, -1);
    lld t = powe(M, mod - 2);
    for(int i = 0; i <= M; i++) {
        a[i] = a[i] * t % mod;
    }
    a.resize(siz);
    return a; 
}
int n, m;
int main() {
    // freopen("data.in", "r", stdin);
    cin >> n >> m;
    vector <lld> a(n + 2), b(m + 2);
    for(int i = 0; i <= n; i++) {
        cin >> a[i];
    }
    for(int i = 0; i <= m; i++) {
        cin >> b[i];
    }
    vector <lld> c = conv(a, b);
    for(int i = 0; i <= n + m; i++) {
        cout << c[i] << " ";
    }
    return 0;
}