FFT/NTT学习笔记

complete_binary_tree · 2025-05-29 16:41:31 · 题解

~~更好的阅读体验~~ 暂无更新支持

Update 2025-10-30：添加了一些 FFT 优化~~谁懂，高精除模板卡常~~，跟进了洛谷的折叠框更新~~为什么没有 question 折叠框~~，并修改了一些表述不当的部分。

多项式乘法优化

如何借助复数和分治，将多项式乘法从 O(n^2) 优化到 O(n \log n)？

本文将介绍两种算法：FFT，FNTT。

FFT

快速傅里叶变换（Fast Fourier Transform）是基于傅里叶变换（Fourier Transform），并使用复数优化的算法，可用于处理多项式乘法。

前置知识：复数

虚数单位用 \mathrm i（不能使用 i）表示，满足 \mathrm i^2=-1。

一种错误理解

复数是一种有两个部分的数，包含实部和虚部，通常表示为 z=x+y\mathrm i，其中 \mathrm i 是虚数单位。

复数的模定义为 \sqrt{x^2+y^2}，通常写作 |z|。

单位根若复数 \omega 满足 \omega^n=1，则称 \omega 是 n 次单位根，可以用 \omega_n 表示。

单位根还可以这样算出：\omega_n=\mathrm e^{\cdot \frac{2k\pi\mathrm{i}}{n}}，其中 k \in [0,n) 且 k \in \mathbb Z。

当然，单位根还有一个性质：若 n 为偶数，则 \omega^{\frac n 2}=-1。

复数加法若 z_1=a_1+b_1\mathrm i，z_2=a_2+b_2\mathrm i，则 z_1+z_2=(a_1+a_2)+(b_1+b_2)\mathrm i。

复数减法若 z_1=a_1+b_1\mathrm i，z_2=a_2+b_2\mathrm i，则 z_1-z_2=(a_1-a_2)+(b_1-b_2)\mathrm i。

复数乘法若 z_1=a_1+b_1\mathrm i，z_2=a_2+b_2\mathrm i，则 z_1 \times z_2=(a_1a_2-b_1b_2)+(a_1b_2+b_1a_2)\mathrm i。

可以证明，复数加减法、复数乘法分别满足交换律、结合律，复数乘法与加减法满足分配律。

多项式点值表示法

我们知道，已知一些点对 (x_i,y_i)，其中 \forall i,j \in [1,n],x_i \not= x_j，那么我们可以唯一确定一个 n-1 次多项式。

这不是插值吗？

还真是。

因为 n 个点可以确定一个 n-1 次多项式。

假设这两个要乘起来的多项式分别是 f(x)、g(x)，那么我们可以取一些值 x_i，则 (x_i,f(x_i) \cdot g(x_i)) 这个点肯定在它们俩乘起来的多项式上。

然而，我们要想办法如何快速实现求值和确定多项式呢？

借助单位根

我们取 \omega=\mathrm{e}^{\frac{2\pi\mathrm{i}}{n}}，取这个值的目的是为了 \omega^n=1，即 \omega 是 n 次单位根。

为什么不取其他使 \omega^n=1 的值？

也可以，不过这样取可以更方便求单位根。

而且有一些单位根可能会存在整数 k<n，使得 \omega_n^k=1，这样子 \{\omega_n^0,\omega_n^1,\omega_n^2,\dots,\omega_n^n\} 中会有重复，就不满足点值表示的要求（点值表示中 x 不可重复）。

我们很容易求出 n 对点对 (\omega^i,f(\omega^i))。

那么怎么知道点对求多项式呢？

设有点对 \{(\omega^i,y_i)\}，多项式为 f(x)=\sum_{i=0}^{n-1}a_ix^i。

直接求太困难了，~~伟大的逆傅里叶变换告诉我们~~我们可以构造一个序列 c 满足

c_i = \sum_{j=0}^{n-1}y_j\omega^{-ij}

也就是分别把点 [\omega^{-0i},\omega^{-i},\omega^{-2i},\dots,\omega^{-(n-1)i}] 代入求多项式 \sum_{i=0}^ny_ix^i 的值。

考虑把 y_j=\sum_{k=0}^{n-1}a_k\omega^{jk} 代入式子：

c_i &= \sum_{j=0}^{n-1}y_j\omega^{-ij}\\ &= \sum_{j=0}^{n-1}\left(\sum_{k=0}^{n-1}a_k\omega^{jk}\right)\omega^{-ij}\\ &= \sum_{k=0}^{n-1}\sum_{j=0}^{n-1}a_k\omega^{jk-ij}\\ &= \sum_{k=0}^{n-1}a_k\sum_{j=0}^{n-1}\omega^{j(k-i)} \end{align*}

我们考虑 \sum_{k=0}^{n-1}a_k\sum_{j=0}^{n-1}\omega^{j(k-i)}：

如果 k=i

那么后面的 \omega^{j(k-i)} 项全都变成了 \omega^{0j}，也就是 1。

所以 c_i=a_i\sum_{j=0}^{n-1}1=na_i。
如果 k \not= i

那么 \sum_{j=0}^{n-1}\omega^{j(k-i)} 相当于一个公比为 \omega^j 的等比数列。

套求和公式，\sum_{j=0}^{n-1}\omega^{j(k-i)}=1^{k-i} \times \frac{1-\omega^{n(k-i)}}{1-\omega^{k-i}}。

又 \omega^n=1，即 1^{k-i} \times \frac{1-\omega^{n(k-i)}}{1-\omega^{k-i}}=1^{k-i} \times \frac{1-1}{1-\omega^{k-i}}=0。

所以 a_k\sum_{j=0}^{n-1}\omega^{j(k-i)}=0。

综上，a_i=\frac {c_i} n，即

a_i = \frac1n \cdot \sum_{j=0}^{n-1}y_j\omega^{-ij}

这个式子和刚刚求多项式的式子差不多（只不过 \omega^{ij} 变成了 \omega^{-ij}），所以我们可以利用类似求多项式值的方法求出。

分治优化

不过到目前算法复杂度还是只能达到 O(n^2)，我们考虑分治优化。

现在我们有了两个式子：

y_i=\sum_{j=0}^{n-1}a_j\omega^{ij}

a_i = \frac1n \cdot \sum_{j=0}^{n-1}y_j\omega^{-ij}

然而朴素地计算它们时间复杂度是 O(n^2) 的，无法接受。

所以我们考虑奇偶分治。（先把 \omega^i 替换成 x 不然看的难受）

f(x)&= (a_0+a_2x^2+a_4x^4+\dots)+(a_1x+a_3x^3+a_5x^5+\dots)\\ &= (a_0+a_2x^2+a_4x^4+\dots)+x(a_1+a_3x^2+a_5x^4+\dots) \end{align*}

设 f_1(x)=a_0+a_2x^1+a_4x^2+\dots，f_2(x)=a_1+a_3x^1+a_5x^2+\dots。

那么

f(x)=f_1(x^2)+xf_2(x^2)

由于进行了平方，所以变成 f_1(x^2) 和 f_2(x^2) 之后只要算一半（单位根性质，\{\omega_{n}^{2i}(i \in [1,n])\}=\{\omega_{\frac n2}^{i}(i \in [1,\frac n2])\}，其中 \omega_n=\mathrm{e}^{\mathrm{i}\frac{2\pi}{n}}）。

单位根还有另外一个性质：

上式代入 \omega^i

y_i=f_1(\omega^{2i})+\omega^{i}f_2(\omega^{2i})

代入 \omega^{i+\frac n 2}

y_{i+\frac n 2} &= f_1(\omega^{2(i+\frac n 2)})+\omega^{i + \frac n 2}f_2(\omega^{2(i+\frac n 2)})\\ &= f_1(\omega^{2i})-\omega^{i}f_2(\omega^{2i}) \end{align*}

可以发现这两个式子（几乎）完全一样。

所以由于单位根性质，只要算 y_0,y_1,\dots,y_{\frac n 2-1} 就能 O(1) 得出 y_{\frac n 2},y_{\frac n 2+1},\dots,y_{n-1}。

综上，我们可以对其进行分治。

时间复杂度 T(n)=O(n)+2T(\frac n 2)，即 O(n \log n)。

核心代码

:::info[我是 FFT 模板]

typedef complex<double> comp;
const double pi=acos(-1.0); // cos(pi)=-1

/// @brief 聚 DFT、IDFT 于一体的 FFT 函数。
/// @param n 需要求的多项式长度。
/// @param a 需要求的多项式数组。结果在此处返回。
/// @param type 1 表示 DFT，-1 表示 IDFT
void FastFastTLE(int n, comp*a, int type){
    if(n==1)return; // 常数项直接返回

    comp*a1=new comp[n>>1],*a2=new comp[n>>1];
    for(int i=0;i<n;++i) // 奇偶分类
        if(i&1)a2[i/2]=a[i];
        else a1[i/2]=a[i]; 

    // 递归分治
    FastFastTLE(n>>1, a1, type);
    FastFastTLE(n>>1, a2, type);

    // wmul 就是单位根，w 表示 wmul^i
    comp wmul=comp(cos(2.0*pi/n),sin(2.0*pi/n)*type), //这里也可以使用 exp 函数，效果一样
         w=comp(1.0,0.0);
    for(int i=0;i<(n>>1);++i){
        comp tmp=w*a2[i]; // 复数计算比较慢，提前计算节省时间
        a[i]=a1[i]+tmp,a[i+(n>>1)]=a1[i]-tmp; // 式子
        w*=wmul;
    }

    delete[] a1;
    delete[] a2;
}

:::

FNTT

快速数论变换（Fast Number-Theoretic Transform）是一种基于数论变换（Number-Theoretic Transform），并利用原根进行优化的算法，可以处理带模数的多项式乘法。

在不引起混淆的语境下，FNTT 也可以简称为 NTT。

前置知识：原根

在模 p 意义下，原根 g 是一个满足以下条件的整数：

在 p 为奇质数时，\varphi (p)=p-1，此时它具有以下性质：

如果我们设 n 为偶数，\omega=g^{\frac {p-1}n}，那么有：

:::info[3 的证明]

由于 g^{p-1}\equiv 1\pmod p，即 g^{2 \cdot \frac{p-1}2}-1 \equiv 0\pmod p。

因式分解，(g^{\frac{p-1}2}+1)(g^{\frac{p-1}2}-1) \equiv 0\pmod p。

由于 g^{\frac{p-1}2} \not \equiv 1，所以 g^{\frac{p-1}2} \equiv -1。

所以 \omega^{\frac{n}2}\equiv -1\pmod p。

:::

根据上面三个性质，我们发现：在模 p 意义下，原根可以代替 FFT 中的单位根。

于是我们可以直接套上 FFT，并把单位根替换成原根，然后再加上模数即可。

:::warning[注意]{open}

NTT 对于模数有要求。

对于一个模数 p，它可以处理的最长序列是 2^k，其中 k 为满足 2^k \mid p 的最大整数。

如果 2^k 太大了会导致 \frac{p-1}{2^k} 不是整数，就做不了 NTT 了。

像常见模数 998344353=7 \times 17 \times 2^{23}，所以模 998244353 最大可以做 n=2^{23} 的序列。

对于任意模数 NTT，请见任意模数多项式乘法。 :::

核心代码

:::info[我是 NTT 模板]

typedef long long ll;
ll W[(1<<21)+5],invW[(1<<21)+5]; // （需要提前预处理）这两个数组在下标为 n 的时候分别表示 n 次单位根、n 次单位根的逆元。

/// @brief NTT 函数。
/// @param n 需要求的多项式长度。
/// @param a 需要求的多项式数组。结果在此处返回。
/// @param type 1 表示 DFT，-1 表示 IDFT。
///
/// 基本架构和 FFT 一样，只不过 complex 变成了 long long，还有需要预处理一下单位根（现场求 O(logn) 对常数影响极大）。
void Nearly_Totally_TLE(int n,ll* a,int type){
    if(n==1)return;
    ll*a1=new ll[n>>1],*a2=new ll[n>>1];
    for(int i=0;i<n;++i)if(i&1)a2[i/2]=a[i];else a1[i/2]=a[i];
    Nearly_Totally_TLE(n>>1,a1,type),Nearly_Totally_TLE(n>>1,a2,type);
    ll wmul=(type+1)?W[n]:invW[n],w=1; // 单位根 comp -> ll
    for(int i=0;i<(n>>1);++i)a[i]=(a1[i]+w*a2[i])%mod,a[i+(n>>1)]=(a1[i]-w*a2[i])%mod,w=w*wmul%mod;
    delete[] a1;
    delete[] a2;
}

:::

优化

我们考虑不使用额外空间，也就是只使用初始的 a 数组完成操作。也就是说，每次相当于把 a 拆成左右两半当作 a_1、a_2 数组。

位逆序置换优化

我们发现，每次 DFT/IDFT 递归的时候都要把 a 数组分为左右两半。我们能不能直接求出原数组会变成什么样呢？

本质上，分两半的操作相当于：

假设在处理长度为 2^n 的序列，做到第 i(0\leq i<2^n) 个；
将 i 的低 n 位循环右移。

:::info[为什么是循环右移] 只考虑修改后的低 n-1 位，那么归类的操作相当于把 i 右移一位。

如果考虑 i 的第 n 位，假设数组 a1 代表的下标第 n 位为 0，数组 a2 为 1，那么我们发现最低位为 0 的数被放到了 a1，最低位为 1 的数放到了 a2。 :::

不难看出这是一个翻转操作。

所以我们可以在 FFT 的时候翻转后再求，大大降低了常数。

那么如何求翻转后是什么呢？

考虑 R(x) 为 x 翻转后的数，考虑递推。显然，R(0)=0。

当 x 末位是 1，那么 R(x)=R(x-1) \text{ bitor } 2^{n-1}。

当 x 末位是 0，那么 R(x)=\frac{R(\frac x2)}2。

于是我们就能写出如下代码：

int rev[(1<<21)+5];
inline void reverse(ll*a,int n){
    for(int i=1;i<n;++i)rev[i]=((i&1)?(rev[i-1]|(n>>1)):(rev[i>>1]>>1));
    for(int i=0;i<n;++i)if(i<rev[i])swap(a[i],a[rev[i]]);//避免重复交换 
}

蝶形运算优化

~~其实这个只能算一个小技巧~~

因为在计算的时候求 a_{i} 和 a_{i+\frac n2} 的时候刚好只会用到 a_{i} 和 a_{i+\frac n2}，所以可以直接在原位置上覆写，不用额外数组。

你可以开两个临时变量来保存它们的值，也可以先算出 a_i，那么 a_{i+\frac n2}:=a_i-2\omega^ia_{i+\frac n2}。

我们不需要 IDFT！

还记得我们构造的序列 c 吗：

c_i = \sum_{j=0}^{n-1}y_j\omega^{-ij}

我们偷偷把 \omega^{-ij} 换成 \omega^{ij}（这样就和 DFT 完全一致了）：

c_i = \sum_{j=0}^{n-1}y_j\omega^{ij}

开推：

c_i &= \sum_{j=0}^{n-1}y_j\omega^{ij}\\ &= \sum_{j=0}^{n-1}\left(\sum_{k=0}^{n-1}a_k\omega^{jk}\right)\omega^{ij}\\ &= \sum_{k=0}^{n-1}\sum_{j=0}^{n-1}a_k\omega^{jk+ij}\\ &= \sum_{k=0}^{n-1}a_k\sum_{j=0}^{n-1}\omega^{j(k+i)} \end{align*}

分类讨论：

k\equiv -i\pmod n
有 a_k\sum_{j=0}^{n-1}\omega^{j(k+i)}=na_k。
k \not\equiv -i\pmod n
有
a_k\sum_{j=0}^{n-1}\omega^{j(k+i)}=a_k \times 1 \times \frac{1-\omega^{n(k+i)}}{1-\omega^{k+i}}=0

综上，我们发现求出的多项式只要除掉 n 再把 1 \sim n-1 位翻转即可。于是我们就可以和 IDFT 说 ByeBye 了。

~~实际上这玩意没什么优化反而最后还要 reverse 一遍，跑得更慢了~~

递归常数大，我要写循环

有了上述三个优化，DFT/IDFT 的递归形式就能很好地转变为循环形式了：

:::info[我是 FFT 模板]

typedef complex<double> comp;
const double pi=acos(-1.0); // cos(pi)=-1
int rev[(1<<21)+5];

/// @brief 聚 DFT、IDFT 于一体的 FFT 函数。
/// @param n 需要求的多项式长度。
/// @param a 需要求的多项式数组。结果在此处返回。
/// @param type 1 表示 DFT，-1 表示 IDFT
///
/// 循环实现，常数优化；IDFT 后无需再处理，更加方便。额外空间 O(1)，避免 MLE。 
inline void FastFastTLE(int n, comp*a, int type){
    for(int i=1;i<n;++i)rev[i]=((i&1)?(rev[i-1]|(n>>1)):(rev[i>>1]>>1));
    for(int i=0;i<n;++i)if(i<rev[i])swap(a[i],a[rev[i]]);//避免重复交换 
    for(int len=1;len<=(n>>1);len<<=1){ //当前处理长度 
        comp wmul(cos(2.0*pi/(len<<1)),type*sin(2.0*pi/(len<<1))); //单位根 
        for(int j=0;j<n;j+=(len<<1)){ //枚举每一段 
            comp w(1.0,0.0);
            for(int k=0;k<len;++k){
                comp a1=a[j+k],a2=a[j+len+k]*w; //存储 
                a[j+k]=a1+a2,a[j+len+k]=a1-a2;  //蝴蝶变换优化 
                w*=wmul;
            }
        }
    }
    if(type<0){for(int i=0;i<n;++i)a[i]/=n;} //IDFT 处理 
}

:::

当然，我们还有 NTT。

:::info[我是 NTT 模板]

const int mod=998244353,N=(1<<21)+5,maxn=1<<21;

int a[N],b[N];
int n,m;

inline ll fstp(ll a,ll b){
    ll ans=1,x=a;
    while(b){
        if(b&1)ans=ans*x%mod;
        x=x*x%mod,b>>=1;
    }
    return ans;
}

int rev[N];

void Nearly_Totally_TLE(int n,int* a,int type){
    for(int i=1;i<n;++i)rev[i]=((i&1)?(rev[i-1]|(n>>1)):(rev[i>>1]>>1));
    for(int i=0;i<n;++i)if(i<rev[i])swap(a[i],a[rev[i]]);//避免重复交换 
    for(int len=1;len<=(n>>1);len<<=1){ //当前处理长度 
        ll wmul=((type==1)?fstp(3,(mod-1)/(len<<1)):fstp(fstp(3,mod-2),(mod-1)/(len<<1))); //单位根 
        for(int j=0;j<n;j+=(len<<1)){ //枚举每一段 
            ll w=1;
            for(int k=0;k<len;++k){
                ll a1=a[j+k],a2=a[j+len+k]*w%mod; //存储 
                a[j+k]=(a1+a2)%mod,a[j+len+k]=((a1-a2)%mod+mod)%mod;  //蝴蝶变换优化 
                w=w*wmul%mod;
            }
        }
    }
    if(type<0){
        ll invn=fstp(n,mod-2);
        for(int i=0;i<n;++i)a[i]=1ll*a[i]*invn%mod;
    }
}

:::

此时我们的程序的速度得到了大幅提升，从 4.78s(avg.531ms,max.1.85s) 优化到了 1.51s(avg.168ms,max.580ms)，提升了将近 4 倍！

当然，这肯定还能继续优化。

一次乘法只需要两次 FFT

我们把 F 的系数存到实部，G 的系数存到虚部，那么就可以构造出这样一个多项式：

H(x)=F(x)+\text iG(x)

那么 H^2(x)=(F^2(x)-G^2(x))+2\text i(F(x)\cdot G(x))。

所以说，我们只要求出 H^2(x) 的虚部除以 2 即可。

用了这个优化我们的程序速度又可以快一些了：1.07s(avg.119ms,max.420ms)。

例题

P3803 【模板】多项式乘法（FFT）| P1919 【模板】高精度乘法 | A*B Problem 升级版

本质相同的两道模板题。

只需对两个序列分别做一遍 DFT，然后把点值相乘，最后再做一遍 IDFT 即可。

使用 FFT、NTT 实现均可，使用 FFT 实现还可以只做一遍 DFT。

需要注意的是，最后答案要除以做 DFT/IDFT 的长度而不是多项式项数，做 DFT/IDFT 的长度要一样且是 2 的整数次幂。还有就是做高精度乘法的时候要记得把数组反向。