NTT学习笔记

和\(FFT\)相对应的,把单位根换成了原根,把共轭复数换成了原根的逆元,最后输出的时候记得乘以原\(N\)的逆元即可.

#include <bits/stdc++.h>
using namespace std;

#define LL long long 

const int MAXN = 3 * 1e6 + 10, P = 998244353, G = 3; 

LL a[MAXN], b[MAXN];
int N, M, limit = 1, L, r[MAXN], Gi;

inline LL fastpow(LL a, LL k) {
    LL base = 1;
    while(k) {
        if(k & 1) base = (base * a ) % P;
        a = (a * a) % P;
        k >>= 1;
    }
    return base % P;
}
inline void NTT(LL *A, int type) {
    for (int i = 0; i < limit; i++) { 
        if(i < r[i]) swap(A[i], A[r[i]]);
    }
    for (int mid = 1; mid < limit; mid <<= 1) {  
        LL Wn = fastpow (type == 1 ? G : Gi , (P - 1) / (mid << 1));
        for(int j = 0; j < limit; j += (mid << 1)) {
            LL w = 1;
            for (int k = 0; k < mid; k++, w = (w * Wn) % P) {
                int x = A[j + k], y = (w * A[j + k + mid]) % P;
                A[j + k] = (x + y) % P;
                A[j + k + mid] = (x - y + P) % P;
            }
        }
    }
}

int main () {
    Gi = fastpow (G, P - 2);
    cin >> N >> M;
    for (int i = 0; i <= N; i++) {cin >> a[i]; a[i] = (a[i] + P) % P;}
    for (int i = 0; i <= M; i++) {cin >> b[i]; b[i] = (b[i] + P) % P;}
    while (limit <= N + M) limit <<= 1, L++;
    for (int i = 0; i < limit; i++) r[i] = (r[i >> 1] >> 1) | ((i & 1) << (L - 1));  
    NTT (a, 1); NTT (b, 1);    
    for (int i = 0; i < limit; i++) a[i] = (a[i] * b[i]) % P;
    NTT (a, -1); 
    LL inv = fastpow (limit, P - 2);
    for (int i = 0; i <= N + M; i++) {
        printf ("%d ", (a[i] * inv) % P);
    }
    return 0;
}
上一篇:任意模数NTT


下一篇:[bzoj3992][SDOI2015]序列统计——离散对数+NTT