【数论】FFT 的原理及 NTT
1. 前言
说起 FFT,本人曾于大概一年前写过一篇相关的文章,但限于本人语文水平、理解程度等问题后来便废弃掉了,于是现在重新写一篇理解较为透彻的文章作为补漏。具体来说,我会采取与之前不同的介绍顺序和角度,用较易理解的语言表述。
2. 多项式基础
2.1 定义
\(\circ\) 定义 \(n\) 次多项式 \(F(x)=x^n+x^{n-1}+\cdots+x+1\),您也可以把它理解成函数,表示多项式和函数时有时将后面的括号省去留下大写字母。
\(\circ\) 定义 \(\deg{F}\) 表示多项式的度,即最高次项的次数,有 \(\deg{F}=n\)
\(\circ\) 定义 \(F[i]\) 表示多项式 \(F(x)\) 第 \(i\) 次项的系数,则有 \(F(x)=\sum_{i=0}^{n}F[i]x^i\)
2.2 卷积
令 \(\oplus\) 为某种运算,有多项式 \(A(x)\) 和 \(B(x)\),则它们在运算 \(\oplus\) 下的卷积
\(C(x)=A(x)B(x) \Rightarrow C[k]=\sum\limits_{i\oplus j=k}A[i]\times B[j]\)
显然,多项式乘法就是加法卷积(类比我们之前学高精度的竖式计算可知),其中,显然有 \(\deg{C}=\deg{A}+\deg{B}\)。
3. FFT 的引入
3.1 DFT 和 IDFT
了解完一些基础知识后,现在我们来进入正题:如何计算多项式乘积?
我们日常计算两个多项式的乘积的时候通常会把它们中的每一项分别相乘,再合并同类项,如:
\((2x^2+x+1)(x^2+3x+2)=2x^4+6x^3+4x^2+x^3+3x^2+2x+x^2+3x+2=2x^4+7x^3+8x^2+5x+2\)
是不是头很晕?(,这样算实在是太慢了!我们考虑优化它。换一种简单的表示,令 \(A(x)=2x^2+x+1\),\(B(x)=x^2+3x+2\),则它们的乘积 \(C(x)=2x^4+7x^3+8x^2+5x+2\)。试试看能不能想到点什么?我们代入几个值到 \(A\),\(B\),\(C\) 看看:
代入 \(x=0\),得 \(A(x)=1\),\(B(x)=2\),\(C(x)=2\)
代入 \(x=1\),得 \(A(x)=4\),\(B(x)=6\),\(C(x)=24\)
代入 \(x=2\),得 \(A(x)=11\),\(B(x)=12\),\(C(x)=132\)
显然,无论 \(x\) 取何值,都有 \(A(x)B(x)=C(x)\)。前面说过,一个多项式 \(F(x)\) 也可以理解成一个函数,那么对于每一个 \(x\) 的取值,都有一个对应的 \(F(x)\),这些点 \((x,F(x))\) 我们称之为 \(F(x)\) 的点值表示。由初中学过的待定系数法可知,\(n+1\) 个点便能确定一个 \(n\) 次函数(多项式),于是我们又能把点值表示转化为系数表示。
我们想,如果我们能把 \(A(x)\) 和 \(B(x)\) 快速地转化为点值表示,在 \(O(n)\) 时间内求出 \(C(x)\) 的点值表示,再快速地将其转化为系数表示,岂不美哉?
于是我们称将系数表示转化为点值表示的算法为 DFT(离散傅里叶变换,Discrete Fourier Transform),将点值表示转化为系数表示的算法为 IDFT (离散傅里叶逆变换,Inverse Discrete Fourier Transform),DFT 的过程我们会优化成 FFT(快速傅里叶变换,Fast Fourier Transform),IDFT 的过程我们会优化成 IFFT(快速傅里叶逆变换,Inverse Fast Fourier Transform)。IDFT 这个算法我们会直接在后文 FFT 部分引出,下面我们讲讲 DFT。
3.2 DFT 的优化
在后面的文章中,始终假定 \(\deg{A}+\deg{B}=\deg{C}=n-1\)且 \(n=2^k\)
(\(k\) 为任意自然数)。
也就是说我们需要找到 \(x\) 的 \(n\) 个取值(我们尽量取 \(2\) 的幂是为了方便后面的运算)并快速计算出对应的 \(A(x)\) 和 \(B(x)\) 的值。问题就出在这,我们不可能枚举 \(n\) 个取值并计算,这样的复杂度是 \(O(n^2)\) 的,落入窠臼,还不如直接暴力卷积呢。于是我们有一个思考方向:怎样取值能使计算量减少呢?如果这个函数是有对称性的,我们枚举一个取值,根据对称性不就可以直接得出对称点的所在的函数值了嘛!于是翻翻百度,找到了函数的奇偶性这个知识点:
简单来说,若一个函数 \(F(x)\),满足 \(F(x)=F(-x)\),则称这个函数为 偶函数;若满足 \(F(x)=-F(-x)\),则称这个函数为 奇函数。可以看出,偶函数关于 y 轴 对称,奇函数关于 原点 对称,两个偶函数的和(或积)是一个偶函数,两个奇函数的和是奇函数,积是偶函数。例如 \(F(x)=x^2\) 是偶函数,\(F(x)=2x^4+3x^2\) 是偶函数,\(F(x)=x^3\) 是个奇函数。
于是我们要求 \(n\) 个 \(x\) 对应的 \(A(x)\) 的值,考虑将 \(A(x)\) 化成以下式子:
也就是说,我们将每一项按次数的奇偶提出来,化成最后的样子,我们发现两边括号里都是偶函数,我们用字母表示出来,左边的式子为 \(A1(x^2)\),右边的式子为 \(A2(x^2)\),则有 \(A(x)=A1(x^2)+x A2(x^2)\),那我们根据偶函数的对称性可以得出 \(A(-x)=A1(x^2)-x A2(x^2)\)。
哇塞!我们这样只用算一半的取值就行了欸,而且 \(A1\) 和 \(A2\) 都变成原来的一半规模,我们把这两个函数再继续这样递归下去,那总时间复杂度就是 \(T(n)=2T(\frac{n}{2})=O(n\log n)\) 了!.................诶?感觉哪里怪怪的?
我们要利用偶函数的对称性进行计算,那么我们取的这些 x 值则作为一对对相反数存在,即便在缩小规模后也是,但是,\(A(x)=A1(x^2)+x\times A2(x^2)\) 中,\(A1\) 和 \(A2\) 的 \(x\) 值都是平方啊, 怎么会是相反数呢?
当然可以,什么样的一组数不断平方后每两个仍是相反数呢?聪明的科学家们把目光投向了一个神奇的领域——复数。
我们有 \(1^2=1\),\(i^2=-1\)。那么对于 \(n=4\) 的时候,我们的 \(x\) 值可以取 \(1\)、\(-1\)、\(i\) 和 \(-i\),递归回溯过程如下:
(1 -1 i -i) → (1 -1) → (1)
具体地,找出这些确切的取值实质上就是求解 \(x^n=1\) 的所有根,于是我们引入了一个有趣的东西—— \(n\) 次单位根。
3.3 \(n\) 次单位根的性质
我们知道,对于一个复数 \(z=a+bi\),其乘法运算为模长相乘,幅角相加,既然求解复数意义下的 \(x^n=1\) 即 \(z^n=1\),那么容易发现其模长始终为 \(1\) 且幅角为 \(\frac{360°}{n}\),我们称这样的 \(z\) 为 \(n\) 次单位根,记作 \(\omega_n\)。那么 \(\omega_n\)、\(\omega_n^2\)、\(\cdots\)、\(\omega_n^{n-1}\)、\(\omega_n^n(=\omega_n^0))\) 正好将复平面的单位圆平分成了 \(n\) 等份,如图:
\(\circ\) 引理:\(\omega_n=\cos\theta+i\sin\theta=e^{i\theta}\)
后半部分是欧拉公式,我们证明下前半部分
因为模长为 \(1\),幅角为 \(\theta\),显然其横坐标为 \(\cos\theta\),纵坐标为 \(\sin\theta\),代入到复平面内就得到 \(\omega_n=\cos\theta+i\sin\theta\)。对于 \(\omega_n^k:1\le k\le n\),有 \(\theta=\frac{2\pi k}{n}\),则有 $$\omega_n^k=\cos\frac{2\pi k}{n}+i\sin\frac{2\pi k}{n}=e^{i\frac{2\pi k}{n}}$$
\(\circ\) 引理:\(\omega_{2n}^{2k}=\omega_n^k\)(折半引理)
证明:
\(\circ\) 引理:\(\omega_n^{k+\frac{n}{2}}=-\omega_n^k:1\le k\le \frac{n}{2}\)(消去引理)
证明:
其中 \(\omega_n^{\frac{n}{2}}\) 相当于 \(\omega_n^0\) 绕逆时针旋转 \(180°\),结合图像可知 \(\omega_n^{\frac{n}{2}}=-1\)。所以根据引理可知 \(\omega_n^{k+\frac{n}{2}}\) 和 \(\omega_n^k\) 互为相反数。
所以我们回到原来的递归式,\(A(x)=A1(x^2)+x A2(x^2)\),\(A(-x)=A1(x^2)-x A2(x^2)\),代入 \(x=\omega_n^k(1\le k\le \frac{n}{2})\),得
这样我们就能够正常地进行递归啦!这样加速后的 DFT 我们就叫它 FFT。
4. FFT 的实现
4.1 从递归到递推
根据上面我们可以得出下面的递归过程:
- 将整个多项式按次数的奇偶性分成两个多项式,并对这两个多项式重复划分直到只有一项。
- 回溯并合并。
由此我们可以写出这样的代码:
void FFT(Complex *A,int n){ // Complex 是手写的复数类,其中包含实部和虚部两个元素
if(n==1)return;
int mid=n>>1;
Complex A1[N],A2[N];
for(int i=0;i<mid;++i){
A1[i]=A[i<<1];
A2[i]=A[(i<<1)|1];
}
FFT(A1,mid),FFT(A2,mid);
Complex wn=Complex(cos(PI/mid),sin(PI/mid)),w=Complex(1,0);//初始化单位根
for(int i=0;i<mid;++i){// 进行合并
A[i]=A1[i]+w*A2[i];
A[i+mid]=A1[i]-w*A2[i];
w=w*wn;
}
return;
}
可尽管是 \(O(n\log n)\) 的复杂度,递归带来的堆栈以及动态开空间等还是吃不消,能不能变递归为迭代呢?
我们发现,变递归为迭代的瓶颈在于,如何快速地将多项式划分,得到最终的序列,然后逐个部分合并?形象地说,递归的思路是个正的二叉树,而迭代的思路是个倒的二叉树,我们的问题是,如何得到迭代的初始状态?
考虑找一下规律,以下是 \(n=8\) 时每一项的编号序列以及划分后的序列:
0 1 2 3 4 5 6 7 -> 0 4 2 6 1 3 5 7
看不太出来?换成二进制试试:
000 001 010 011 100 110 101 111 -> 000 100 010 110 001 011 101 111
发现了什么?划分后的序列每一项实质上就是原序列每一项的二进制反转!
考虑递推求出每一个数的二进制反转(也称位逆序置换,Bit Reverse),定义 \(rev_i\) 表示 \(i\) 的二进制反转后的数,有 \(rev_{rev_i}=i\) 。假设 \(rev_{\frac{i}{2}}\) 已知,则我们知道 \(i\) 二进制除第一位的反转,我们只需确定反转后的最高位即可,将 \(rev_{\frac{i}{2}}\) 右移一位(因为 \(\frac{i}{2}\) 反转后最高位移到最低位多占了一位),补上最高位,就得到 \(i\) 的二进制反转了,代码:
for(int i=0;i<n;++i)rev[i]=(rev[i>>1]>>1)|((i&1)<<len-1); // len 是 n 的二进制位数
for(int i=0;i<n;++i)if(i<rev[i])swap(A[i],A[rev[i]]);// 处理
现在考虑递归的问题,我们要做的就是:
- 枚举长度。
- 对每一块进行计算合并。
得到如下代码:
void FFT(Complex A,int n){
for(int mid=1;mid<n;mid<<=1){ // 枚举每次合并区间长的一半
Complex wn=Complex(cos(PI/mid),sin(PI/mid));
for(int len=mid<<1,now=0;now<lim;now+=len){ // now 为每一块的起始位置
Complex w=Complex(1,0);
for(int j=0;j<mid;++j,w=w*wn){
Complex x=A[now+j],y=w*A[now+j+mid]; // 此时的 A 为未合并的状态
A[now+j]=x+y;
A[now+j+mid]=x-y;
}
}
}
return;
}
这个 \(x\) 和 \(y\) 加减的过程有个有意思的名字:“蝴蝶变换“。
4.2 代码流程梳理
首先,我们进行二进制反转的操作,然后通过蝴蝶变换逐个合并,最终得到点值表示。如图:
4.3 IFFT
我们解决了系数转点值的过程,现在我们要来考虑点值转系数的过程了。
已知 \(A(\omega^i_n)\times B(\omega^i_n)=C(\omega^i_n)=\sum\limits_{j=0}^{n-1}\omega^{ij}_nc_j\),求 \(c_i=?\)。
发现可以写成矩阵相乘的形式(此处将 \(C(\omega^i_n)\)) 简写成 \(C_i\))
我们要求 \(c_i\) 的矩阵就将左右两边的等式乘上中间矩阵的逆就行了,我们有以下结论:
浅证一下:
考虑 \(\sum\limits_{i=0}^{n-1}\omega^{i(k-j)}_n\) 的取值,当 \(k=j\) 时
当 \(k\not= {j}\) 时,代入等比数列求和
所以得
由此,我们可以直接用上面 FFT 的函数做 IFFT 运算,只需把 \(\omega_n\) 变为 \(\omega_n^{-1}\) ,最后除以 \(n\) 就行了。
4.4 完整实现
\(\circ\) 例题:P3803 【模板】多项式乘法(FFT)
提交记录(注意 double 的精度,输出时要四舍五入)
\(\bullet\) 三步变两步优化
考虑复数 \(z=a+bi\),有 \(z^2=(a+bi)^2=a^2-b^2+2abi\)
发现了什么?我们把输入的多项式 \(B\) 放在 \(A\) 的虚部上,求出 \(A^2\) 就行了。
提交记录(快了 1s)
5. 升级:从 FFT 到 NTT
FTT 虽好,但其单位根为浮点数的限制太大,容易引发精度问题。考虑到单位根只是一个用来加速运算的东西,且大多的计数问题都需要取模,那我们能否在模意义下找到一个单位根的替代品呢?
答案是还真有(废话,不然我写这个干嘛),我们叫它 NTT(快速数论变换,Number Theory Transform)。
5.1 原根
对于 \(a\in \mathbb{Z}\),\(p\in\mathbb{N}^+\),\(\gcd(a,p)=1\),满足 \(a^n\equiv 1\pmod p\) 的最小的正整数 \(n\),我们称它为“\(a\) 模 \(p\) 的阶”,记作 \(\delta_p(a)\) 或 \(\operatorname{ord}_p(a)\)。显然由欧拉定理可知 \(n\) 的上界为 \(\varphi(p)\)。
若 \(g\in\mathbb{Z}\) 满足 \(\delta_p(g)=\varphi(p)\) ,即达到上界,那我们称 \(g\) 是模 \(p\) 意义下的原根。下面所述的 \(g\) 都代表原根。
\(\circ\) 引理:\(g^1,g^2,\cdots,g^{\delta_p(g)}\) 在模 \(p\) 意义下两两不同余。
证明:
若存在 \(i\not=j\),且 \(g^i\equiv g^j\pmod p\),那么有 \(g^{|i-j|}\equiv 1\pmod p\),可 \(1\leq|i-j|< \delta_p(g)\),与阶的最小性矛盾,故引理成立。
这也说明了为什么能用原根的原因之一:我们使用单位根的原因之一,就是因为它们在 \(n\) 次下取值各不相同。但 \(n\) 显然不一定等于 \(p-1\),那么我们肯定要想办法利用原根构造一组取值。
5.2 原根的使用
对于质数 \(p=qn+1(n=2^k)\) (一般给的质数都满足这个),则其原根 \(g\) 满足 \(\delta_p(g)=qn\),那么我们令 \(g_n=g^q=g^{\frac{p-1}{n}}\) 就可以构造出等价于 \(\omega_n\) 的值啦。
诶诶诶等等,为什么啊?
要证明 \(g_n\) 和 \(\omega_n\) 等价,我们得要证明 \(g_n\) 跟 \(\omega_n\) 有一样的性质,那么回顾一下,\(\omega_n\) 有哪些性质呢?
\(\circ\) 折半引理:\(\omega_{2n}^{2k}=\omega_n^k\)
证明:
所以,\(g_{2n}^{2k}\equiv g_n^k\pmod p\)。
\(\circ\) \(\omega_n^{\frac{n}{2}}=-1\)
证明:
我们有
若 \(g^{\frac{p-1}{2}}\equiv 1\pmod p\),则与阶的最小性矛盾,那么只能有 \(g^{\frac{p-1}{2}}\equiv -1\pmod p\)。
所以,\(g_n^{\frac{n}{2}}\equiv -1\pmod p\)。
\(\circ\) 消去引理: \(\omega_n^{k+\frac{n}{2}}=-\omega_n^k\)
证明:
所以,\(g_n^{k+\frac{n}{2}}\equiv-g_n^k\pmod p\)
由此,我们只需要把所有 \(\omega_n\) 都换成模 \(p\) 意义下的 \(g_n\) 就好啦!
5.3 NTT 的实现
还是模板题 Luogu P3803 【模板】多项式乘法(FFT)
提交记录(注意逆变换中需要逆元)
比三步变两步优化后的 FFT 还要快一些,因为我们无需常数极大的复数计算,以及一些三角函数计算,实际上还可以再优化。
5.4 附:原根表
一般常用的质数及其原根有:
6. 后记 & 参考资料
简单地写了一下,就当作复习和学习新的东西,如果有什么错漏还请提出,我会及时修改。
\(\circ\) FFT 参考资料:
繁凡さん:【学习笔记】超简单的快速傅里叶变换(FFT)(含全套证明)
Bilibili 视频:快速傅里叶变换(FFT)——有史以来最巧妙的算法?
Bilibili 视频:具体学习并实现快速傅里叶变换(FFT)