FFT和NTT学习笔记

算法导论

概述

目的

以\(O(nlg_n)\)的时间复杂度计算多项式乘法

多项式的表达

系数表达: \(\{a_0, a_1, ..., a_{n-1}\}\)
点值表达: \(\{(x_0,y_0), (x_1,y_1), ..., (x_{n-1},y_{n-1})\}\)
- 插值多项式唯一性: \(x\)各不相同得到唯一的多项式

策略

graph TD
A[A,B的系数表达] -- 普通乘法n^2 --> B[C的系数表达]
A[A,B的系数表达] -- 扩展成2n并求值nlgn --> A2[A,B在2n个单位复数根的点值]
A2[A,B在2n个单位复数根的点值] -- 点值乘法n --> B2[C在2n个单位复数根的点值]
B2[C在2n个单位复数根的点值] -- 插值nlgn --> B[C的系数表达]

关于单位复数根

复平面

复数平面（complex plane）是用水平的实轴与垂直的虚轴建立起来的复数的几何表示。

一个复数的实部用沿着 x-轴的位移表示，虚部用沿着 y-轴的位移表示。

复数乘法：模长相乘，幅角相加（幅角：向量与x轴正向夹角）

单位复数根

感性理解:

将复平面上的单位圆从 x 轴起分成 n 等分，从 x 轴起标号 0..n-1

那么由 “模长相乘，幅角相加” 可得，标号为 1 的那个复数的 \(k\) 次方，就是标号为 \(k\) 的

所以将这些数记为 \(\omega_n^0, \omega_n^1, \dots, \omega_n^{n-1}\)

理性证明:

\(n\)次单位复数根 是满足 \(\omega^n=1\) 的复数 \(\omega\), 恰好有 \(n\) 个

由 \(e^{i\theta}=\cos(\theta)+i\sin(\theta)\)

得 \(e^{2\pi i}=\cos(2\pi)+i\sin(2\pi)=1\)

所以对于 \(k=0,1,...n-1\) , 这些根是 \(e^{2\pi i k/n}\)
有复数乘法可得这 \(n\) 个单位复数根均匀分布在以复平面原点为圆心的单位半径的圆周上
\(\omega_n=e^{2\pi i/n}\) 称为 主\(n\)次单位根，其他单位根都是 \(\omega_n\) 的幂次，并有 \(\omega_n^{j+k}=\omega_n^{(j+k) \mod n}\)

单位复数根基本性质

(消去引理)

对任何整数 \(以及n\geq0,k\geq0,以及d\geq0\) , \(\omega_{dn}^{dk}=\omega_n^k\)
感性理解：原来分成 \(n\) 等分，取第 \(k\) 个，和分成 \(dn\) 等分，取第 \(dk\) 个当然一样
证明：根据指数形式定义式

(折半引理)

感性理解：

\(\omega_n^{k + n/2} = -\omega_n^k\)

复平面单位圆上关于原点对称

理性证明：

如果 \(n>0\) 为偶数，那么 \(n\) 个 \(n\) 次单位复数根的平方的集合就是 \(n/2\) 个 \(n/2\) 次单位复数根的集合
证明：\((\omega_n^{k+n/2})^2=\omega_n^{2k+n}=\omega_n^{2k}\omega_n^n=(\omega_n^k)^2=\omega_{n/2}^k\)
- 也可以根据 \(\omega_n^{n/2}=-1\) 得到 \(\omega_n^{k+n/2}=-\omega_n^k\)
由此可见他可以递归让子问题的规模缩小一半

(求和引理)

\[\sum_{j=0}^{n-1}(\omega_n^k)^j =
\begin{cases}
n, & k = 0 \\
0, & k \neq 0 \\
\end{cases}
\]

感性理解：

\(n\) 是 \(2\) 的次幂

\(k=0\) 显然

\(k \neq 0\) 时

首先一个偶数 \(n\) 的等分点的和是 \(0\) (对称点), 那么对于 \(n=2^k\) 显然适用

假如 \(k\) 是 \(n\) 的因数, 那么 \(k\) 也是 \(2\) 的因数, 它会取一个等分的循环取若干次, 和是 \(0\)

否则互质, 那么它会取遍这 \(n\) 个点, 和也是 \(0\)

证明：等比数列求和 \(\sum_{j=0}^{n-1}(\omega_n^k)^j=\frac{1*(1-\omega_n^{kn})}{1-\omega_n^k}=0\)
因为要求 \(k\) 不能被 \(n\) 整除，保证分母不为 \(0\)

概述

系数转点值: 求出 \(n\) 个单位根的点值 \((y_0, y_1, \dots, y_{n-1})\)

点值转系数: 将这 \(n\) 个点值作为一个新的多项式 \(B(x)\) 的系数, 用单位根的倒数求一次点值 \((z_0, z_1, \dots, z_{n-1})\)

展开可得

\[\begin{align*}
z_k &= \sum_{i = 0}^{n - 1} y_i(\omega_n^{-k})^i \\
&= \sum_{i = 0}^{n - 1}(\sum_{j = 0}^{n - 1} a_j(\omega_n^i)^j)(\omega_n^{-k})^i \\
&= \sum_{j = 0}^{n - 1}a_j(\sum_{i = 0}^{n - 1}(\omega_n^{j - k})^i)
\end{align*}
\]

根据之前的求和引理可得: \(z_k = n \cdot a_k\)

于是系数 \(a_k = z_k / n\)

也就是说, 这两个过程都可以通过求点值来完成,

所以要完成两个多项式的相乘, 只需要先求一遍他们的点值, 点值相乘, 在转成系数就行了

DFT

点值向量\(y=(y_0, y_1, ..., y_{n-1})\)就是系数向量\(a=(a_0, a_1, ..., a_{n-1})\)的离散傅里叶变换(DFT), 也记为\(y=DFT_n(a)\)

FFT

快速傅里叶变换(FFT)，利用单位复数根的特殊性质，在\(O(nlg_n)\)的时间内计算出\(DFT_n(a)\)

分治策略，采用\(A(x)\)中偶数下标和奇数下标的系数，分别定义两个次数界为\(n/2\)的多项式

\(A_0(x) = a0 + a2*x + ... + a_{n-2}*x^{n/2-1}\)
\(A_1(x) = a1 + a3*x + ... + a_{n-1}*x^{n/2-1}\)

\(A(x) = A_0(x^2) + x \cdot A_1(x^2)\)

带入单位根可以发现

\[\begin{align*}
A(\omega_n^k) &= A_0(\omega_n^{2k}) + \omega_n^kA_1(\omega_n^{2k}) \\
&= A_0(\omega_{\frac{n}{2}}^{k}) + \omega_n^kA_1(\omega_{\frac{n}{2}}^{k})
\end{align*}
\]

\[\begin{align*}
A(\omega_n^{k + \frac{n}{2}})
&= A_0(\omega_n^{2k + n}) + \omega_n^{k + \frac{n}{2}}A_1(\omega_n^{2k + n}) \\
&= A_0(\omega_{\frac{n}{2}}^{k} \times \omega_n^n) + \omega_n^{k + \frac{n}{2}} A_1(\omega_{\frac{n}{2}}^{k} \times \omega_n^n) \\
&= A_0(\omega_{\frac{n}{2}}^{k}) - \omega_n^kA_1(\omega_{\frac{n}{2}}^{k})
\end{align*}
\]

这两个关于远点对称的单位根的点值只是相差一个符号, 所以计算 \(A_0(\omega_{\frac{n}{2}}^{k}), \omega_n^kA_1(\omega_{\frac{n}{2}}^{k})\) 就可得到两个单位根的取值

因为应用了\(\omega_n^k\)的正负数形式，所以称其为旋转因子

可得到以下伪代码，时间复杂度\(O(nlg_n)\)

RECURSIVE_FFT(a) // 计算DFTn(a)
    n = a.length
    if n == 1
        return a // a * (x^0)
    a[0] = (a(0), a(2), ..., a(n-2))
    a[1] = (a(1), a(3), ..., a(n-1))
    y[0] = RECURSIVE_FFT(a[0])
    y[1] = RECURSIVE_FFT(a[1])
    for k = 0 to n/2-1
        y[k] = y[0][k] + omega[n][k] * y[1][k]
        y[k+n/2] = y[0][k] - omega[n][k] * y[1][k]
    return y;

IDFT

通过点值得到系数

IFFT

在概述中说明了, 只需要将点值作为系数的到新的多项式, 再带入单位根的倒数, 得到这个多项式的点值

然后点值 \(/n\) 即可得到系数

高效FFT实现

画出系数的递归树可得叶子的标号的规律:

蝴蝶操作中

graph LR
id1[y0k] --> A
id2[+] --> ret1[y0k+omega*y1k]
id3[y1k] --> B
id4[-] --> ret2[y0k-omega*y1k]
A --> id2[+]
A --> id4[-]
B --> id2[+]
B --> id4[-]
ret1[y0k + omega*y1k] --> ret3[y k]
ret2[y0k - omega*y1k] --> ret4[y k+n/2]

可以观察到\(y_k^{[1]}\)的位置其实就是\(y_{k+n/2}\)的位置，相当于两个位置上的\(y\)值进行蝴蝶操作并只对这两个位置进行修改

观察下图递归FFT输入的向量树

graph TD
id1[a0,a1,a2,a3,a4,a5,a6,a7] --- id2[a0,a2,a4,a6]
id1[a0,a1,a2,a3,a4,a5,a6,a7] --- id3[a1,a3,a5,a7]
id2[a0,a2,a4,a6] --- id4[a0,a4]
id2[a0,a2,a4,a6] --- id5[a2,a6]
id3[a1,a3,a5,a7] --- id6[a1,a5]
id3[a1,a3,a5,a7] --- id7[a3,a7]
id4[a0,a4] --- id8[a0,000]
id4[a0,a4] --- id9[a4,100]
id5[a2,a6] --- id10[a2,010]
id5[a2,a6] --- id11[a6,110]
id6[a1,a5] --- id12[a1,001]
id6[a1,a5] --- id13[a5,101]
id7[a3,a7] --- id14[a3,011]
id7[a3,a7] --- id15[a7,111]

其叶子有这样的规律：

记\(rev(x)\)为\(x\)的二进制表示的倒串所形成的\(lg_n\)位的数,则\(rev(x)\)位上为\(a_x\)

因为每次的叶子是一样的( \(n\) 相同), 所以可以预处理, 并且可以 \(O(n)\) 递推 rev[]

注意预处理叶子位置不是扫到 \(len / 2\)

不妨从下往上模拟递归的过程

方便起见(省去二维的数组), 将每一层中各多项式的各点值标号, 将递归的写成这样

{
    FFT(a, n / 2);
    FFT(a + n / 2, n / 2);
    for (int i = 0; i < n / 2; ++ i) // 枚举 \omega_n^i
    {
        buf[i] = a[i] + a[i + m];
        buf[i + m] = a[i] - \omega_{n / 2}^{-i} * a[i + m];
    }
}

转成非递归, 现将原系数变成叶子的状态

for (int l = 2; l <= len; l <<= 1)
    {
	int ll = l >> 1;
	for (int i = 0; i < len; i += l)
	{
	    for (int j = 0; j < l / 2; ++ j)
	    {
		CPD x = ome[len / l * j];
		if (opt) x = inv[len / l * j];
		buf[i + j] = a[i + j] + x * a[i + j + ll];
		buf[i + j + ll] = a[i + j] - x * a[i + j + ll];
	    }
	}
	for (int i = 0; i < len; ++ i) a[i] = buf[i];
    }

注意这里的单位根可以利用引理对应到 \(n\) 次单位根上

卡常提示:

然后注意不要预处理 ome[], inv[], 循环的时候直接乘上第一个单位根即可
手写复数, 并且不要用构造函数(常数小了一半)
2 * PI / l 可以改为 PI / mid

// P1919 A*B(FFT)
//#pragma GCC optimize(2)
#include <iostream>
#include <cstdio>
#include <cmath>
#include <cstring>
#include <algorithm>
using namespace std;
struct Complex
{
    double x, y;
    Complex operator + (const Complex &b) const { return (Complex){x + b.x, y + b.y}; }
    Complex operator - (const Complex &b) const { return (Complex){x - b.x, y - b.y}; }
    Complex operator * (const Complex &b) const { return (Complex){x * b.x - y * b.y, x * b.y + y * b.x}; }
} ;
const int MAXN = (1 << 21) + 3;
char ch[MAXN];
int lena, lenb, n, dgt;
Complex a[MAXN], b[MAXN];
void read(Complex * a, int & len)
{
    scanf("%s", ch + 1); len = strlen(ch + 1);
    for (int i = 0; i < len; ++ i) a[i].x = ch[len - i] - '0';
}
int rev[MAXN];
void init(int n, int dgt)
{
    for (int i = 0; i < n; ++ i)
	rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (dgt - 1));
}
const double PI = acos(-1.0);
void FFT(Complex * a, int len, int opt)
{
    for (int i = 0; i < len; ++ i)
	if (i < rev[i]) swap(a[i], a[rev[i]]);
    for (int mid = 1; mid < len; mid <<= 1)
    {
	Complex omen = (Complex){cos(PI / mid), opt * sin(PI / mid)} ;
	for (int i = 0; i < len; i += (mid << 1))
	{
	    Complex ome = (Complex){1, 0} ;
	    for (int j = 0; j < mid; ++ j, ome = ome * omen)
	    {
		Complex t = ome * a[i + j + mid];
		a[i + j + mid] = a[i + j] - t;
		a[i + j] = a[i + j] + t;
	    }
	}
    }
}
int ret[MAXN];
/*
20191212
0724~0757~0839
a * b FFT
 */
int main()
{
    read(a, lena); read(b, lenb);
    n = 1; dgt = 0;
    while (n < lena + lenb) n <<= 1, ++ dgt;
    init(n, dgt);
    FFT(a, n, 1);
    FFT(b, n, 1);
    for (int i = 0; i < n; ++ i) a[i] = a[i] * b[i];
    FFT(a, n, -1);
    for (int i = 0; i < n; ++ i)
    {
	ret[i] += (int)(a[i].x / n + 0.5);
	ret[i + 1] += ret[i] / 10;
	ret[i] %= 10;
    }
    while (!ret[n - 1] && n > 0) -- n;
    for (int i = n - 1; i >= 0; -- i) printf("%d", ret[i]);
//    printf("#%d\n", cnt);
    return 0;
}
/*
76543210
76543210
6543
21
 */

总结

FFT 利用了单位根的以下性质

from http://blog.miskcoo.com/2015/04/polynomial-multiplication-and-fast-fourier-transform#i-17

首先来看 FFT 中能在 O(nlogn) 时间内变换用到了单位根 ω 的什么性质 ...

自己解释一遍:

单位根 \(\omega_n^0, \dots, \omega_n^{n-1}\) 是单位圆的等分点(不同), 所以有了 "消去引理"
\(\omega_n^{\frac{n}{2}+k}=-\omega_n^k, \omega_n^2=\omega_{\frac{n}{2}}\)
- 根据奇偶系数拆分的表达式, 这一层的 \(m\) 个多项式都要带入 \(n\) 个值, 可以通过下一层 \(2m\) 个多项式带入 \(n/2\) 个值求得, 由此可知 \(\log\) 层
\(\sum_{j=0}^{n-1}(\omega_n^k)^j =
\begin{cases}
n, & k = 0 \\
0, & k \neq 0 \\
\end{cases}\) 使得点值转系数的逆变换可以同样方式完成

其中第一点是最基本的, 他推出了后面的性质

NTT

在取一些模数的意义下, 原根具有上述性质

区别在于原根的手动计算不同

void NTT(LL * a, int len, int sgn)
{
    for (int i = 0; i < len; ++ i)
	if (i < rev[i]) swap(a[i], a[rev[i]]);
    for (int mid = 1; mid < len; mid <<= 1)
    {
	for (int i = 0; i < len; i += (mid << 1))
	{
	    for (int j = 0; j < mid; ++ j)
	    {
		LL t = (sgn == 1 ? g[len / mid / 2 * j] : ig[len / mid / 2 * j]) * a[i + j + mid] % MOD;
		a[i + j + mid] = (a[i + j] + MOD - t) % MOD;
		a[i + j] = (a[i + j] + t) % MOD;
	    }
	}
    }
}

invn = inv(n);
g[0] = ig[0] = 1;
g[1] = fpow(3, (MOD - 1) / n); ig[1] = inv(g[1]);
for (int i = 2; i < n; ++ i)
    g[i] = g[i - 1] * g[1] % MOD, ig[i] = ig[i - 1] * ig[1] % MOD;