FFT与NTT的模板

网上相关博客不少，这里给自己留个带点注释的模板，以后要是忘了作提醒用。

以洛谷3803多项式乘法裸题为例。

 

FFT：

#include <cstdio>
#include <cmath>
#include <cctype>
#include <algorithm>
#define ri readint()
#define gc getchar()

int readint() {
    int x = 0, s = 1, c = gc;
    while (c <= 32)    c = gc;
    if (c == '-')    s = -1, c = gc;
    for (; isdigit(c); c = gc)    x = x * 10 + c - 48;
    return x * s;
}

const int maxn = 4 * 1e6 + 10;
const double PI = acos(-1.0);

struct Complex {
    double x, y;
    Complex(double a = 0, double b = 0):x(a), y(b){}
};
Complex operator + (Complex A, Complex B) { return Complex(A.x + B.x, A.y + B.y); }
Complex operator - (Complex A, Complex B) { return Complex(A.x - B.x, A.y - B.y); }
Complex operator * (Complex A, Complex B) { return Complex(A.x * B.x - A.y * B.y, A.x * B.y + A.y * B.x); }

Complex a[maxn], b[maxn];
int n, m;
int r[maxn], l, limit = 1;

void fft(Complex *A, int type) {
    for (int i = 0; i < limit; i++)
        if (i < r[i])
            std::swap(A[i], A[r[i]]);
    //迭代方式模拟递归写法，需要理解递归是怎么做的才能看懂这个
    for (int mid = 1; mid < limit; mid <<= 1) {
        //本来单位根是2*PI/len，这里len替换成2*mid，2就约掉了
        Complex Wn(cos(PI / mid), type * sin(PI / mid));
        for (int R = mid << 1, j = 0; j < limit; j += R) {
            Complex w(1, 0);//单位根的k次幂
            for (int k = 0; k < mid; k++, w = w * Wn) {
                //蝴蝶变换
                Complex x = A[j+k], y = w * A[j+k+mid];
                A[j+k] = x + y;
                A[j+k+mid] = x - y;
            }
        }
    }
}

int main() {
    n = ri, m = ri;
    for (int i = 0; i <= n; i++)
        a[i].x = ri;
    for (int i = 0; i <= m; i++)
        b[i].x = ri;

    while (limit <= n + m) {//长度变为2^l
        limit <<= 1;
        l++;
    }
    for (int i = 0; i < limit; i++)//二进制镜像
        r[i] = (r[i>>1] >> 1) | ((i&1) << (l-1));
    fft(a, 1);
    fft(b, 1);
    for (int i = 0; i < limit; i++)
        a[i] = a[i] * b[i];
    fft(a, -1);
    for (int i = 0; i <= n + m; i++)
        printf("%d ", (int)(a[i].x / limit + 0.5));
    return 0;
}

 

　NTT是用模域取代了复数域，性质相同只是换了单位根，所以板子基本相同。我这两个相比NTT确实比FFT快一点的：

#include <bits/stdc++.h>
#define ll long long
#define ri readll()
#define gc getchar()
#define rep(i, a, b) for (int i = a; i <= b; i++)
using namespace std;

const int P = 998244353, G = 3, Gi = 332748118, maxn = 4 * 1e6 + 5;
//P的原根为3,3%P的逆元为332748118
//原根意味着：3^(P-1) % P = 1,其中P-1是3%P的阶，本应是φ(P)，这里恰好为大素数
ll n, m;
ll a[maxn], b[maxn];
int limit = 1, l, r[maxn];

ll readll() {
    ll x = 0ll, s = 1ll;
    char c = gc;
    while (c <= 32)    c = gc;
    if (c == '-')    s = -1ll, c = gc;
    for (; isdigit(c); c = gc)    x = x * 10 + c - 48;
    return x * s;
}

ll ksm(ll a, ll b, int mod) {
    ll res = 1ll;
    for (; b; b >>= 1) {
        if (b & 1)    res = res * a % mod;
        a = a * a % mod;
    }
    return res;
}

void NTT(ll *A, int flag) {
    rep(i, 0, limit)
    if (i < r[i])
        swap(A[i], A[r[i]]);

    for (int mid = 1; mid < limit; mid <<= 1) {
        //如果是变换则单位根为3^[(P-1)/(len)] % P,逆变换则用逆元
        ll Wn = ksm(flag ? G : Gi, (P-1) / (mid*2), P);
        for (int R = mid << 1, j = 0; j < limit; j += R) {
            ll w = 1ll;
            for (int k = 0; k < mid; k++, w = w * Wn % P) {
                ll x = A[j+k], y = A[j+k+mid] * w % P;
                A[j+k] = (x + y) % P;
                A[j+k+mid] = (x - y + P) % P;
            }
        }
    }
}

int main() {
    n = ri, m = ri;
    rep(i, 0, n)    a[i] = (ri + P) % P;
    rep(i, 0, m)    b[i] = (ri + P) % P;

    while (limit < n + m + 1) {
        limit <<= 1;
        l++;
    }
    rep(i, 0, limit)    r[i] = (r[i>>1] >> 1) | ((i & 1) << (l - 1));
    NTT(a, 1);    NTT(b, 1);
    rep(i, 0, limit)    a[i] = a[i] * b[i] % P;
    NTT(a, 0);

    ll inv = ksm(limit, P - 2, P);//最后变换回来要乘长度的逆元
    rep(i, 0, n + m)    printf("%lld ", a[i] * inv % P);

    return 0;
}