浅谈FFT与NTT

Posted by K0u1e on March 28, 2020

本来觉得明明去年暑假里就学过了现在却忘得干干净净了真是太丢人了,看了zzy的ppt以后发现自己根本就没有学会过,这里%%%zzy。

离散傅里叶变换

设$a,b$分别为两个多项式的系数数组,类型为复数,度数界分别为$n$和$m$。现在要求多项式系数数组$c=a \cdot b$。

这里不妨把$a,b$的度数界统一定为不小于$n+m-1$的$2$的幂次$n$。

首先把$a$和$b$两数组转为点值表达式,然后根据$c_i=a_i\cdot b_i$求出$c$数组的点值表达式,再由$c$数组的点值表达式求$c$数组的系数表达式。

也就是让列向量$a$和$b$分别乘上一个矩阵$H$得到$A,B$两个列向量,即点值表达形式,由$A,B$求出列向量$C$后再由$C$乘上$H^{-1}$得到$c$。大概是下面这个过程。

\[H \cdot a = \begin{bmatrix} x_0^0 & x_0^1 & x_0^2 & \cdots & x_0^{n-1} \\ x_1^0 & x_1^1 & x_1^2 & \cdots & x_1^{n-1} \\ x_2^0 & x_2^1 & x_2^2 & \cdots & x_2^{n-1} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ x_{n-1}^0 & x_{n-1}^1 & x_{n-1}^2 & \cdots & x_{n-1}^{n-1} \end{bmatrix} \cdot \begin{bmatrix} a_0\\ a_1\\ a_2\\ \vdots \\ a_{n-1}\\ \end{bmatrix} = \begin{bmatrix} A(x_0)\\ A(x_1)\\ A(x_2)\\ \vdots \\ A(x_{n-1})\\ \end{bmatrix} =A\]

这里引入复数$w_n=\cos\frac{2\pi}{n}+i\sin\frac{2\pi}{n}$的次幂作为$T$的组成元素,即令$x_i=w_n^i$。这样的$n$个$x$两两不同并且具有下面两个性质:

  1. $w_n^{i+n}=w_n^i$
  2. $w_{2n}^{2i}=w_n^i$

所以我们的矩阵$H$的各个元素$H_{ij}=w_n^{ij}$,$H^{-1} _ {ij} = \frac{1}{n} w_n^{-ij} = H _ {(n-i)j}$,完整的过程可以表示为:

\[H \cdot a= \begin{bmatrix} w_n^0 & w_n^0 & w_n^0 & \cdots & w_n^0\\ w_n^0 & w_n^1 & w_n^2 & \cdots & w_n^{n-1}\\ w_n^0 & w_n^2 & w_n^4 & \cdots & w_n^{n-2}\\ \vdots & \vdots & \vdots & \ddots & \vdots \\ w_n^0 & w_n^{n-1} & w_n^{n-2} & \cdots & w_n^1\\ \end{bmatrix} \cdot \begin{bmatrix} a_0\\ a_1\\ a_2\\ \vdots \\ a_{n-1}\\ \end{bmatrix} = \begin{bmatrix} A(w_n^0)\\ A(w_n^1)\\ A(w_n^2)\\ \vdots \\ A(w_n^{n-1})\\ \end{bmatrix} =A\] \[H^{-1} \cdot A= \begin{bmatrix} w_n^0 & w_n^0 & w_n^0 & \cdots & w_n^0\\ w_n^0 & w_n^{n-1} & w_n^{n-2} & \cdots & w_n^1\\ w_n^0 & w_n^{n-2} & w_n^{n-4} & \cdots & w_n^2\\ \vdots & \vdots & \vdots & \ddots & \vdots \\ w_n^0 & w_n^{1} & w_n^{2} & \cdots & w_n^{n-1}\\ \end{bmatrix} \cdot \begin{bmatrix} A(w_n^0)\\ A(w_n^1)\\ A(w_n^2)\\ \vdots \\ A(w_n^{n-1})\\ \end{bmatrix} = \begin{bmatrix} a_0\\ a_1\\ a_2\\ \vdots \\ a_{n-1}\\ \end{bmatrix} =a\]

关于$H^{-1}$的正确性的证明:

\[\sum_{k=0}^{n-1}H_{ik}\cdot H^{-1}_{kj}=\sum_{k=0}^{n-1}w_n^{k(i-j)}\\ = \begin{cases} 0\quad i\neq j\\ n \quad i=j \end{cases}\]

快速傅里叶变换

刚刚那个思路的朴素实现复杂度依旧是$O(n^2)$,但我们可以考虑分治将$a$与$A$的相互转化优化到$O(nlogn)$。

我们将$A(x)$根据奇偶项拆分成$A(x)=A_0(x^2)+xA_1(x^2)$。

正常来讲我们依旧要把每个$w_n^i$代入,但是我们会发现这些值的$w_n^{2i}=w_{\frac{n}{2}}^i$有只有$\frac{n}{2}$个,有一半是重复的,也就是说我们只要代入一半即可。这样我们就把问题规模缩小了一半,递归处理完两个规模减半的子问题$A_0$与$A_1$以后再$O(n)$合并成$A$即可。

快速数论变换

有时候多项式相乘之后系数会变得奇大无比,需要对某个质数$p$取模,这时候普通的FFT就满足不了我们的要求了。一种方法是使用快速数论变换,也就是NTT。

在FFT中我们找到复数$w_n$来作为单位根,主要是因为$w_n^n=1$导致$w_n^{i+n}=1$,使得有效的$x^2$只有$\frac{n}{2}$个。在NTT中,如果我们能找到一个$g_n^n \equiv 1 \mod p$,那么我们就可以用类似的方法来。

原根:令$p$为一正整数。对于一个正整数$g$,若对任意一个与$p$互质的正整数$a$,均存在整数$k$,使$a\equiv g^k\mod p$, 则我们称$g$为模$p$意义下的一个原根。

设$p=m\cdot 2^k+1$($m$为奇数),取$n \mid 2^k$,$g_n=g^{\frac{p-1}{n}}$。

由费马小定理可知$g^{p-1}\equiv 1 \mod p$,那么显然$g_n^{n}\equiv 1 \mod p$。于是我们就有了一个等价于$w_n$的单位根来用类似的方式实现NTT。

任意模数NTT

上面的$NTT$只能处理模数满足$p=m\cdot 2^k+1$,且$n \mid 2^k$的情况,如果$p$为任意质数,就不是那么友好了。

三模数

由于一次NTT处理完以后新的多项式一般每个系数都不会超过$n\cdot \max(a_i)^2\le 10^5 \cdot 10^9 \cdot 10^9$,我们考虑用三个互质的满足要求的模数$p_1,p_2,p_3$且$p_1p_2p_3 > 10^{23}$,做三次完整的NTT,对每个值都求一个通余方程组:

\[\begin{cases} ans\equiv a_1 \mod p_1 \\ ans\equiv a_2 \mod p_2 \\ ans\equiv a_3 \mod p_3 \\ \end{cases}\]

由于直接用中国剩余定理会爆$longlong$所以可以先用中国剩余定理算出一个$a_4 \equiv ans \mod p1p2$。

不妨令$ans=a_5p_1p_2+a_4$。

由于$a_5p_1p_2+a_4 \equiv a_3 \mod p_3$,那么显然有$a_5\equiv (a_4-a_3)p_1^{-1}p_2^{-1} \mod p3$。

即$a_5=kp_3+a_6$,$a_6=((a_3-a_4)p_1^{-1}p_2^{-1} \mod p3)$。

代入上面的等式中得到$ans=a_6p_1p_2+a_4$。