快速傅里叶变换（FFT）基础-526互联

本文是对 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\).

多项式多点求值是对系数向量的线性变换，变换后的结果是点值向量. 这个变换矩阵称作范德蒙德矩阵：

\[\begin{bmatrix} 1&x_0&x_0^2&\cdots&x_0^{n-1}\\ 1&x_1&x_1^2&\cdots&x_1^{n-1}\\ 1&x_2&x_2^2&\cdots&x_2^{n-1}\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ 1&x_{n-1}&x_{n-1}^2&\cdots&x_{n-1}^{n-1} \end{bmatrix} \begin{bmatrix} a_0\\a_1\\a_2\\\vdots\\a_{n-1} \end{bmatrix}= \begin{bmatrix} y_0\\y_1\\y_2\\\vdots\\y_{n-1} \end{bmatrix}. \]

DFT 可以看作给系数向量左乘一个特殊的范德蒙德矩阵

\[\begin{bmatrix} 1&1&1&\cdots&1\\ 1&w&w^2&\cdots&w^{n-1}\\ 1&w^2&w^4&\cdots&w^{2(n-1)}\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ 1&w^{n-1}&w^{2(n-1)}&\cdots&w^{(n-1)(n-1)} \end{bmatrix}= \sum_{i,j}w^{ij}e_{ij}. \]

于是，IDFT 相当于给点值向量左乘逆矩阵. 可以用拉格朗日插值求得这个逆矩阵，我们这里直接给出逆矩阵的形式（证明见下文）：

\[\dfrac1n\left[\begin{matrix} 1&1&1&\cdots&1\\ 1&w^{-1}&w^{-2}&\cdots&w^{-(n-1)}\\ 1&w^{-2}&w^{-4}&\cdots&w^{-2(n-1)}\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ 1&w^{-(n-1)}&w^{-2(n-1)}&\cdots&w^{-(n-1)(n-1)} \end{matrix}\right]= \dfrac1n\sum_{i,j}w^{-ij}e_{ij}. \]

其中 \(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\) 的逆就行.

我们还注意到（提升注意力）上面的矩阵恰好等于

\[\newcommand\udots{\mathinner{\kern1mu\raise1pt{.}\kern2mu\raise4pt{.}\kern2mu\raise7pt{\Rule{0pt}{7pt}{0pt}.}\kern1mu}} \begin{bmatrix} 1&&&&&\\ &&&&&1\\ &&&&1&\\ &&&1&&\\ &&\udots&&&\\ &1&&&&\\ \end{bmatrix}\begin{bmatrix} 1&1&1&\cdots&1\\ 1&w&w^2&\cdots&w^{n-1}\\ 1&w^2&w^4&\cdots&w^{2(n-1)}\\ \vdots&\vdots&\vdots&\ddots&\vdots\\ 1&w^{n-1}&w^{2(n-1)}&\cdots&w^{(n-1)(n-1)} \end{bmatrix} \]

（这很容易验证，只须注意到同一行上第 \(i\) 个与第 \(n-i\) 个分量的指数之和总是 \(n\) 的整数倍）. 于是有另一种常见的 IDFT 实现方法：做一遍正常的 DFT，把系数向量乘以 \(n^{-1}\)，再把 \(1\sim n-1\) 翻转. 这样写略微简洁一些.

最后我们做一个简单的乘法，验证它确实是逆矩阵. 需要预先准备一个特殊的等比数列求和

\[\sum_{i=0}^{n-1}(w^k)^i=\begin{cases}0,&n\not\mid k,\\n,&n\mid k.\end{cases} \]

证明是朴素的，只须注意到

\[\left(1-w^k\right)\sum_{i=0}^{n-1}(w^k)^i=1-w^{kn}=0. \]

于是

\[\begin{aligned} (\sum_{i,j}w^{ij}e_{ij})(\sum_{k,l}w^{-kl}e_{kl})&=\sum_{i,j,k,l}w^{ij-kl}e_{ij}e_{kl} \\&=\sum_{i,j,l}w^{(i-l)j}e_{il} \\&=\sum_{i,j,l}w^{(i-l)j}e_{il} \\&=\sum_{i,l}e_{il}\sum_{j=0}^{m-1}(w^{i-l})^j \\&=\sum_{i,l}e_{il}[i=l]n \\&=nI. \end{aligned} \]

后排提示：这里记号 \(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);
    }
}