快速傅里叶变换(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\).

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

\[\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);
    }
}
posted @ 2023-08-19 21:31  abcc!  阅读(112)  评论(0编辑  收藏  举报