【多项式】快速数论变换（NTT）

· · 算法·理论

上一篇：【多项式】两次 FFT 优化
例题：P3803【模板】多项式乘法（FFT）

Part 0：从 FFT 出发

在上一篇我们已经指出了，因为 FFT 是用 double 类型进行计算的，所以难免会有精度误差。那有没有针对整数的零误差计算方法呢？

有，这就是 NTT，即快速数论变换（\text{Number Theory Transform}）。我们将从原根开始，介绍什么是 NTT。

Part 1：原根

我们仅对原根作一个粗略介绍，更详细的推导等可以前往洛谷模板题「原根」查看。

原根的定义：对于正整数 p，若有正整数 a 满足：

存在正整数 x 使得 a^x\equiv 1\pmod p，且 x 的最小值为 \varphi(p)。

那么我们就称 a 是 p 的原根。

举个例子，3 是 5 的原根（\varphi(5)=4），因为：

(x) (1) (2) (3) (4)
(3^x\bmod 5) (3) (4) (2) (1)

但 4 不是 5 的原根，因为 4^2\equiv 1\pmod 5，x 的最小值是 2 而不是 4。

并且我们有一个结论，质数一定有原根，比如说 998244353 的原根就是 3（一般原根不止一个），证明可到模板题去看，此处略。

那么这个原根有啥用呢？诶，我们可以拿它替换掉复数，因为它具有和原根类似的性质。

假设模数 p 的原根是 a，那么 a^{\varphi(p)}\equiv 1\pmod p，设 n\mid \varphi(p)。

在复数中，我们有 U_n^k=-U_n^{k+\frac{n}{2}}，即 U_n^{\frac{n}2}=-1。

因为同样有 a^{\frac{\varphi(p)}{2}}\equiv -1\pmod p，所以 U_n^{\frac{n}2} 等效于 a^{\frac{\varphi(p)}{2}}，那么我们就可以把 U_n 等效于 a^{\frac{\varphi(p)}{n}}，就可以通过将 FFT 中的 U_n 换成 a^{\frac{\varphi(p)}{n}} 来实现 NTT。

但是如果 \frac{\varphi(p)}{n} 不是个整数怎么办？这简单，因为 n 在计算时取的是 2 的幂，且一般小于等于 2^{21}。那么只要找到一个质数 p 满足 2^{21}\mid p-1 就好了，有这种数吗？

有，比如 469762049,998244353,1004535809 这三个就是，并且 3 都是它们的原根。

Part 2：代码实现

上面把 NTT 的核心讲了，但可能还是不知道要怎么做，我们接着看。

首先来看看怎么把 FFT 函数改成 NTT 的函数，这是 FFT 的源代码：

void FFT(Complex F[], int n, int type){
    for(int i = 1; i < n; ++i) if(i < Rev[i]) swap(F[i], F[Rev[i]]);
    for(int i = 1; i < n; i <<= 1){
        Complex Unit = {cos(Pi / i), type * sin(Pi / i)};
        for(int j = 0; j < n; j += (i << 1)){
            Complex Temp = {1.0, 0.0};
            for(int k = j; k < j + i; ++k){
                Complex F0 = F[k], F1 = F[k + i];
                F[k] = F0 + Temp * F1;
                F[k + i] = F0 - Temp * F1;
                Temp *= Unit;
            }
        }
    }
}

只需把三行 Complex 的改掉就好了。首先是 Unit，因为 U_n 等价于 a^{\frac{\varphi(p)}{n}}，那我们就快速幂一下，计算 a 的 \frac{\varphi(p)}{n} 次方就好了，此处 p=998244353,a=3。如果 type = -1 就计算 a^{-1} 的 \frac{\varphi(p)}{n} 次方，a^{-1}\equiv 332748118\pmod {998244353}。

然后第二行改成 ll Temp = 1 即可，第三行同理改成 ll 类型就好了。哦，函数传参的类型也要改改，注意要对计算结果取模。最后代码长这样：

void NTT(ll F[], int n, int type){
    for(int i = 1; i < n; ++i) if(i < Rev[i]) swap(F[i], F[Rev[i]]);
    for(int i = 1; i < n; i <<= 1){
        ll Unit = FastPow((type == 1) ? 3 : 332748118, 998244352 / (i << 1));
        for(int j = 0; j < n; j += (i << 1)){
            ll Temp = 1;
            for(int k = j; k < j + i; ++k){
                ll F0 = F[k], F1 = Temp * F[k + i] % Mod;
                F[k] = (F0 + F1) % Mod, F[k + i] = (F0 - F1) % Mod; 
                Temp = Temp * Unit % Mod;
            }
        }
    }
}

然后主函数基本与 FFT 的一致，除了改输入输出外，还有那个除以 len 的，要改成乘上 len 的逆元。最重要的一点：注意取模。总代码在此查看：Code（稍微卡了常的代码）。

Part 3：算法分析

NTT 因为是用整数型变量进行运算的，所以不存在精度误差，这点优于 FFT。但 FFT 适用面更广，可以将系数是小数的多项式相乘。

另外，NTT 对模数有一定要求，必须满足 2^{21}\mid (p-1)，如果题目的模数不满足这一要求，就需要任意模数 NTT 来解决，效率反倒不如 FFT。

还有，上篇我们提到 FFT 有三次变两次的优化，那么 NTT 有没有呢？很可惜，没有。虽然我们可以利用二次剩余等效地表示出虚数单位 i，但是 NTT 没有实部和虚部，计算后实部和虚部是混在一起的，无法分开，所以无法借助复数的巧妙性质三次变两次。