快速傅里叶变换(FFT)基础
本文是对 FFT 和 NTT 原理及实现的介绍,包含所有必要的证明. 阅读本文需要具备一点基本的代数知识.
给定 \(n\) 次多项式 \(F(x)\) 和 \(m\) 次多项式 \(G(x)\),现在要求它们的卷积 \(H(x)=F(x)G(x)\). 朴素的暴力实现复杂度为 \(O(nm)\),而 FFT 或 NTT 可以(在一定的精度范围内或模意义下)做到 \(O((n+m)\log(n+m))\).
算法介绍
设 \(R\) 是环, \(R[x]\) 为其上的多项式环,\(F,G\in R[x]\) 为次数小于 \(2^{n-1}\) 的多项式. 离散傅里叶变换(DFT)将 \(F\) 和 \(G\) 由系数表示法变换成点值表示法,由于在点值表示法下多项式的卷积就是对应位置的乘积,这就能求出 \(F*G\) 的点值表示;然后离散傅里叶逆变换(IDFT) 将点值表示法转换回系数表示法. 朴素的求值和插值都是 \(O(n^2)\) 复杂度,但这里我们对求值的点没有要求,只要恰当选取,就有可能应用特殊的方法加快速度.
注意到(提升注意力),只要选取 \(R^\times\) 上的一个 \(2^k\) 阶循环群的全体元素,就能满足足够多的性质(后文给出详细说明). 这要求 \(R^\times\) 上存在阶为 \(2^n\) 的元,很遗憾,最常见的 \(\mathbb Z\) 和 \(\mathbb R\) 上没有这样的元. 一种解决方法是将范围扩展到 \(\mathbb C\),这样全体 \(2^n\) 次单位根在复数乘法下恰好构成循环群. 从技术角度来说,这引入了大量浮点数运算(还全是复数),时间常数巨大. 快速数论变换(NTT)给出了在模 \(p\) 意义下求整系数多项式卷积的方法. 由于 \(\mathbb F_p^\times\) 天然是 \(p-1\) 阶循环群,只要选取的 \(p\) 恰好满足 \(2^n\mid p-1\)(即 \(p-1\) 具有足够多的因子 \(2\)),就能构造出 \(2^n\) 阶循环群:设 \(p\) 的原根为 \(g\)(即 \(g\) 的阶为 \(p-1\)),则 \(g^{(p-1)/2^n}\) 的阶就是 \(2^n\). 如果初始系数很小(或题目就是要求模 \(p\) 意义下的答案),NTT 给出的就是真正的答案,否则可以分别做不同模数的 NTT 再用中国剩余定理合并(MTT).
\(3\) 是质数 \(998244353\) 和 \(1004535809\) 的原根.
在下文中,记 \(w\) 为 \(R^\times\) 中一个阶为 \(2^n\) 的元素,则 \(w^i\ (0\le i<2^n)\) 形成了选取的循环群. 而且,\(w^{2^{n-k}\cdot i}\ (0\le i<2^k)\) 形成了 \(2^k\) 阶的循环子群.
DFT
DFT 的实现是基于倍增的(如果多项式次数不是 \(2\) 的整数幂,要对高次项补 \(0\)).
对于 \(2^k\) 次多项式的系数表示 \(F=\left<a_0,a_1,a_2,\cdots,a_{2^k-1}\right>\),我们的目标是求出全部 \(w^{2^{n-k}\cdot i}\ (0\le i<2^k)\) 的点值.
将 \(F\) 按奇次项和偶次项拆分,得到两个新的 \(2^{k-1}\) 次多项式 \(G=\left<a_0,a_2,\cdots,a_{2^k-2}\right>\) 和 \(H=\left<a_1,a_3,\cdots,a_{2^k-1}\right>\). 我们有 \(F(x)=G(x^2)+xH(x^2)\).
假设我们已经递归地求出 \(G\) 和 \(H\) 在 \(w^{2^{n-k+1}\cdot i}\ (0\le i<2^{k-1})\) 处的点值(注意循环群的大小减半,这保证了分治的复杂度正确). 现在要合并得到 \(F\) 的 \(2^k\) 个点值. 对于每个 \(i\) 计算 \(F(w^{2^{n-k}\cdot i})=G(w^{2^{n-k+1}\cdot i})+w^{2^{n-k}\cdot i}H(w^{2^{n-k+1}\cdot i})\),由于\(F\) 和 \(G\) 自变量的平方恰好使循环群的大小减半,所有需要的点值都已经求过,这便可以合并了.
时间代价与最经典的分治法一样,有 \(T(2^n)=2T(2^{n-1})+2^n\),解得复杂度为 \(\Theta(n2^n)\).
IDFT
前排提示:下文的 \(n\) 含义与上文不同,表示某个 \(2\) 的幂次,\(w\) 是一个 \(n\) 阶循环群的生成元. 矩阵单位 \(e_{ij}\) 的第 \(i\) 行第 \(j\) 列(下标从 \(0\) 开始)为 \(1\),其他位置为 \(0\). 容易验证,\(e_{ij}e_{jk}=e_{ik}\);当 \(j\ne j'\) 时 \(e_{ij}e_{j'k}=0\).
多项式多点求值是对系数向量的线性变换,变换后的结果是点值向量. 这个变换矩阵称作范德蒙德矩阵:
DFT 可以看作给系数向量左乘一个特殊的范德蒙德矩阵
于是,IDFT 相当于给点值向量左乘逆矩阵. 可以用拉格朗日插值求得这个逆矩阵,我们这里直接给出逆矩阵的形式(证明见下文):
其中 \(w^{-1}\) 的阶和 \(w\) 相等,因此也是一个 \(n\) 阶循环群的生成元. 于是我们可以用 \(w^{-1}\) 再次做 DFT,最后把系数向量乘以 \(n^{-1}\),就完成了 IDFT.
注:严谨地讲,在一般的环 \(R\) 上,实际上是乘以 \(n\cdot1=\underbrace{1+1+\cdots+1}_{n个}\) 的逆,其中 \(1\) 为乘法单位元. 其实,因为 \(n\) 总是 \(2\) 的整数幂,只要能求 \(2=1+1\) 的逆就行.
我们还注意到(提升注意力)上面的矩阵恰好等于
(这很容易验证,只须注意到同一行上第 \(i\) 个与第 \(n-i\) 个分量的指数之和总是 \(n\) 的整数倍). 于是有另一种常见的 IDFT 实现方法:做一遍正常的 DFT,把系数向量乘以 \(n^{-1}\),再把 \(1\sim n-1\) 翻转. 这样写略微简洁一些.
最后我们做一个简单的乘法,验证它确实是逆矩阵. 需要预先准备一个特殊的等比数列求和
证明是朴素的,只须注意到
于是
后排提示:这里记号 \(n\) 的含义恢复,现在多项式 \(F\) 和 \(G\) 的次数小于 \(2^{n-1}\).
递推实现
DFT 是递归的过程,但我们可以自底向下地向上合并,实现非递归且原地(直接把系数数组改成点值数组)的 DFT.
首先对于长度为 \(1\) 的多项式 \(F(x)=a_0\),其系数表示与点值表示相同,就不用管了,只需要从长度为 \(2\) 开始合并. 现在的问题是怎么做奇偶次分组. 注意到(提升注意力),\(i\) 次项在最下面一层的位置恰好是把 \(i\) 看作 \(n\) 位 \(2\) 进制数后的位逆序. 这是因为,每次判断奇偶相当于检查二进制的最低位,而分到左右两边相当于设置新下标的最高位. 因此,最后的结果相当于整个下标的二进制位逆序.
为了做到完全原地,还要注意一个细节. 当合并两个 \(2^{k-1}\) 次多项式 \(G\) 和 \(H\) 时,我们要枚举 \(0\le i<2^k\). 此时 \(G\) 和 \(H\) 在内存里是连续相邻的,它们最终要被 \(F(x)\) 的点值表示覆盖. 不妨设目前计算到的 \(i<2^{k-1}\) (在左半部分),我们有 \(F(w^{2^{n-k}\cdot i})=G(w^{2^{n-k+1}\cdot i})+w^{2^{n-k}\cdot i}H(w^{2^{n-k+1}\cdot i})\) 即 \(f_i=g_i+w^{2^{n-k}\cdot i}h_i\). 这会覆盖掉 \(g_i\),而未来计算 \(f_{i+2^{k-1}}=g_i+(w^{2^{n-1}})w^{2^{n-k}\cdot i}h_i\) 还要用到 \(g_i\)(因为分治后循环的周期减半了),所以只须枚举一半范围的 \(i\),同时把 \(f_i\) 和 \(f_{i+2^{k-1}}\) 计算出来并覆盖掉即可. \(h_i\) 前的系数可以先处理出 \(w^{2^{n-k}}\) 并不断自乘. 另外,后一半范围 \(h_i\) 的额外系数 \(w^{2^{n-1}}\) 恰好等于 \(-1\)(乘法单位元的加法逆),这是循环群的性质决定的(因为阶为 \(2\)).
三次变两次优化
求多项式乘法总共要做三次 DFT. 对于 \(\mathbb C\) 上的 FFT 算法,我们通常做实(整)系数多项式卷积时,前两次 DFT 的虚部都被浪费了,能不能把 \(F\) 和 \(G\) 分别存在某个 \(H\) 的实虚部里呢?设 \(H:=F+G\mathrm i\),注意到(提升注意力)\(H^2=(F+G\mathrm i)^2=(F^2-G^2)+2FG\mathrm i\),我们用两次 DFT 即可求出 \(H\) 的平方的系数表示,每一项虚部的一半就是答案.
三次变两次优化后的 FFT 比 NTT 稍快.
样例代码
template<bool INV=false, typename T>
void DFT(T F[], int n){ // precondition: n==2^k
for(int i=0; i<n; i++) rev[i]=rev[i>>1]>>1|(i&1)<<(n-1);
for(int i=0; i<n; i++) if(i<rev[i]) swap(F[i], F[rev[i]]);
for(int len=2; len<=n; len<<=1){
const T w0=exp(complex<double>(0, 2*M_PI/len))/* g^(mod-1)/len */;
const int mid=len>>1;
for(int i=0; i<n; i+=len){
T w(1);
for(int j=i; j<i+mid; j++){
const T x=F[j], y=w*F[j+mid];
F[j]=x+y, F[j+mid]=x-y;
w*=w0;
}
}
}
if(INV){
const auto invn=T(1)/n;
for(int i=0;i<n;i++) F[i]*=invn;
reverse(F+1, F+n);
}
}