「学习笔记」FFT及NTT入门知识

前言

快速傅里叶变换（( ext{Fast Fourier Transform,FFT}) ）是一种能在(O(n log n))的时间内完成多项式乘法的算法，在(OI)中的应用很多，是多项式相关内容的基础。下面从头开始介绍( ext{FFT})。

前置技能：弧度制、三角函数、平面向量。

多项式

形如(f(x)=a_0+a_1x+a_2x^2+...+a_nx^n)的式子称为(x)的(n)次多项式。其中(a_0,a_1,...,a_n)称为多项式的系数。

系数表达法

上面定义中的表示就是系数表达法。其系数可看成(n+1)维向量(vec a=(a_0,a_1,...,a_n))。

点值表达法

把多项式看成一个函数，点值表示就用它图像上的(n+1)个不同的点((x_0,y_0),...,(x_n,y_n))来确定这个多项式。多项式有不止一个点值表示，可以证明每个点值表示确定唯一的系数表达多项式。

复数

虚数单位

(i)被称为虚数单位。规定(i=sqrt {-1})。

复平面

复数的平面由(x,y)轴组成。(x)轴称为实轴，(y)轴称为虚轴。平面内的每一个从原点到某个点((a,b))的向量(vec a=(a,b))表示复数(a+bi).

复数的模长：(sqrt {a^2+b^2}).实轴到复数向量的转角( heta)称为幅角。

复数的基本运算

• 复数的加（减）法：((a+bi)+(c+di)=(a+c)+(b+d)i)
• 复数的乘法：((a+bi)(c+di)=(ac-bd)+(bc+ad)i)
• 一个口诀：复数乘法，模长相乘，幅角相加。可以用下面将提到欧拉公式证明。

共轭复数

(a+bi)与(a-bi)互为共轭复数。图像上是关于x轴对称的。

单位根

(n)次单位根是满足(z^n=1)的(n)个复数，它们均分复平面的单位圆。

这些复数满足模长为(1)，幅角的(n)倍是(2pi)的倍数

欧拉公式：(e^{xi}=cos x+i sin x)，其中(e)为自然对数的底数，(i)为虚数单位。

（p.s. 欧拉公式的证明可以使用泰勒展开）

(n)次单位根有n个，为(e^{frac{2pi ki}{n}},kin [0,n-1])

我们记(omega_n=e^{frac{2pi i}{n}},)则(n)次单位根为(omega_n^0,...,omega_n^{n-1})

上面是数学上单位根定义，还有一种更易理解的方法：

单位根均分复平面单位圆，那(w_n)的幅角就是(frac{2pi}{n})，根据三角函数即可得到。

单位根的性质

性质(1)：根据定义得到：(omega_{2n}^{2k}=omega_{n}^{k})（被叫做折半定理，是消去定理的特殊情形）

性质(2)：(omega_{n}^{frac{n}{2}+k}=-omega_n^k)

证明：

(omega_{n}^{frac{n}{2}}=e^{frac{2pi i}{n} frac{n}{2}}=e^{pi i}=cos pi+i sin pi=-1)

(omega_{n}^{frac{n}{2}+k}=omega_{n}^{frac{n}{2}}omega_{n}^{k}=-omega_n^k)

性质(2)图像上理解是转半圈。

Fast Fourier Transform

多项式乘法

系数表达的多项式乘法：(c(x)=a(x)b(x))，则(c(x)=sum_{i=0}^{2n} c_i x^i)

其中(c_i=sum_{j=0}^{n}a_jb_{i-j}).

时间复杂度(O(n^2))

点值表达的多项式乘法：取相同的一组的(x)，(y)值相乘即可

时间复杂度(O(n))

因此多项式乘法的基本思路是先插值得到点值表达，再(O(n))乘，最后求值得到系数表达。

DFT

把(n)次单位根(omega_n^0,...,omega_n^{n-1})带入多项式(A(x)=a_0+a_1x+...+a_{n-1}x^{n-1})，

得到点值向量(vec y=(A(omega_n^0),A(omega_n^1),...,A(omega_n^{n-1})))，

称为系数向量(vec a=(a_0,a_1,...,a_n))的离散傅里叶变换（( ext{Discrete Fourier Transform, DFT})），写作(vec y= ext{DFT}_n(vec a))。

直接求( ext{DFT})是(O(n^2))的。( ext{FFT})的常用算法( ext{Cooley-Tukey})使用分治方法做到(O(nlog n)).

以下讨论基于(n=2^m,m in N^*)，若不足则高位系数补(0).（注意这里(n)是次数界，即最高次项为(x^{n-1})）

考虑点值向量的第(k+1)维（(k<frac{n}{2})）：

(A(omega_{n}^{k})=sum_{i=0}^{n-1}a_i(omega_{n}^{k})^{i}=sum_{i=0}^{n-1}a_iomega_{n}^{ki})

(=sum_{i=0}^{frac{n}{2}-1}a_{2i}omega_{n}^{2ki}+sum_{i=0}^{frac{n}{2}-1}a_{2i+1}omega_{n}^{2ki+k})

(=sum_{i=0}^{frac{n}{2}-1}a_{2i}omega_{n}^{2ki}+omega_{n}^{k}sum_{i=0}^{frac{n}{2}-1}a_{2i+1}omega_{n}^{2ki})

利用性质1：(omega_{2n}^{2k}=omega_{n}^{k})

(A(omega_{n}^{k})=sum_{i=0}^{frac{n}{2}-1}a_{2i}omega_{frac{n}{2}}^{ki}+omega_{n}^{k}sum_{i=0}^{frac{n}{2}-1}a_{2i+1}omega_{frac{n}{2}}^{ki})

利用性质2：(omega_{n}^{frac{n}{2}+k}=-omega_n^k)

可以推出：

(A(omega_{n}^{k+frac{n}{2}})=(-1)^{frac{n}{2}}sum_{i=0}^{frac{n}{2}-1}a_{2i}omega_{frac{n}{2}}^{ki}+(-1)^{frac{n}{2}}omega_{n}^{k+frac{n}{2}}sum_{i=0}^{frac{n}{2}-1}a_{2i+1}omega_{frac{n}{2}}^{ki})

(A(omega_{n}^{k+frac{n}{2}})=sum_{i=0}^{frac{n}{2}-1}a_{2i}omega_{frac{n}{2}}^{ki}-omega_n^ksum_{i=0}^{frac{n}{2}-1}a_{2i+1}omega_{frac{n}{2}}^{ki})

为了简记，偶数、奇数项系数用(A_0,A_1)表示，则上面的式子可以写成以下形式

(A(omega_{n}^{k})=A_0(omega_{frac{n}{2}}^k)+omega_{n}^{k}A_1(omega_{frac{n}{2}}^k))

(A(omega_{n}^{k+frac{n}{2}})=A_0(omega_{frac{n}{2}}^k)-omega_{n}^{k}A_1(omega_{frac{n}{2}}^k))

上面把求和分成(0,2,...,n-2)与(1,3,...,n-1)两部分，把大小为(n)的问题转化成两个规模为(frac{n}{2})的子问题，可以进行分治求解了。

IDFT

求值过程使用离散傅里叶逆变换（( ext{Inverse Discrete Fourier Transform, IDFT})）

结论：只要把( ext{DFT})的(omega_n)都取倒数（共轭复数），最后除以(n)即可。

证明：

设(vec Y=(y_0,y_1,...,y_n))为(vec A = (a_0,a_1,...,a_n))的离散傅里叶变换。

考虑一个向量：(vec C=(c_0,c_1,...,c_n))满足(c_k=sum_{i=0}^{n-1}y_i(omega_n^{-k})^i)

（即(vec C)是多项式(vec Y)在(omega_n^0,omega_n^{-1},...,omega_n^{-(n-1)})处的点值）

将上式展开：

(c_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_{i=0}^{n-1}(sum_{j=0}^{n-1}a_j(omega_n^j)^i)(omega_n^{-k})^i)

(=sum_{i=0}^{n-1}sum_{j=0}^{n-1}a_j(omega_n^j)^i(omega_n^{-k})^i)

(=sum_{i=0}^{n-1}sum_{j=0}^{n-1}a_j(omega_n^{j-k})^i)

(=sum_{j=0}^{n-1}a_j(sum_{i=0}^{n-1}(omega_n^{j-k})^i))

考虑一个前缀和(S(omega_n^k)=1+omega_n^k+(omega_n^k)^2+...+(omega_n^k)^{n-1})。

当((omega_n^k) ot = 1)即(k ot = 0)时，使用等比数列求和方法：

(omega_n^kS(omega_n^k)=omega_n^k+(omega_n^k)^2+(omega_n^k)^3+...+(omega_n^k)^{n})

(omega_n^kS(omega_n^k)-S(omega_n^k)=(omega_n^k)^{n}-1)

(S(omega_n^k)=frac{(omega_n^k)^{n}-1}{omega_n^k-1})

分母不为(0)，分子((omega_n^k)^{n}-1=(omega_n^n)^{k}-1=1^{k}-1=0)

因此(k ot = 0)时，(S(omega_n^k)=0)

(k=0)时，(S(omega_n^k)=n(omega_n^k)^0=n)

继续考虑刚刚那个式子

(c_k=sum_{j=0}^{n-1}a_j(sum_{i=0}^{n-1}(omega_n^{j-k})^i))

只有(j-k=0)时(sum_{i=0}^{n-1}(omega_n^{j-k})^i)才为(n)，否则为(0)。

(c_k=a_kn)

得到结论：(a_k=frac{c_k}{n})。

上面也可以用矩阵证（更加优美），太麻烦了就不写了。

再次总结一下，离散傅里叶逆变换就是先求出多项式在(omega_n^0,omega_n^{-1},...,omega_n^{-(n-1)})处的点值表示，再每一项除以(n)。和上面的分治求值没有区别，因为这些点值性质完全相同。

递归版代码

按照如上所述的方法可以轻松写出一份递归代码。

//Luogu P3803 多项式乘法 - 递归FFT
#include <complex>
#include <cstdio>
using namespace std;

typedef complex<double> comp;

const int N = (1 << 20) + 10 << 1;
const double PI2 = 2.0 * acos(-1.0);

int read() {
	int x = 0; char c = getchar();
	for(; c < '0' || c > '9'; c = getchar()) ;
	for(; c >= '0' && c <= '9'; c = getchar())
		x = x * 10 + (c & 15);
	return x;
}

int n, m;
comp a[N], b[N];

void fft(int n, comp * a, int type) {
	if(n == 1) return ;
	comp a1[n >> 1], a2[n >> 1];
	for(int i = 0; i < n; i += 2)
		a1[i >> 1] = a[i], a2[i >> 1] = a[i + 1];
	fft(n >> 1, a1, type), fft(n >> 1, a2, type);
	comp w(1, 0), wn(cos(PI2 / n), type * sin(PI2 / n));
	for(int i = 0; i < n >> 1; i ++, w *= wn)
		a[i] = a1[i] + w * a2[i], 
		a[i + (n >> 1)] = a1[i] - w * a2[i];
}

int main() {
	n = read(), m = read();
	for(int i = 0; i <= n; i ++) a[i] = read();
	for(int i = 0; i <= m; i ++) b[i] = read();
	
	int lim = 1;
	for(; lim <= n + m; lim <<= 1) ;
	
	fft(lim, a, 1), fft(lim, b, 1);
	for(int i = 0; i <= lim; i ++) a[i] *= b[i];
	fft(lim, a, -1);
	
	for(int i = 0; i <= n + m; i ++)
		printf("%d ", (int)(0.5 + a[i].real() / lim));
	return 0;
}

迭代优化

通过观察得到：多项式的(i)次项到分治边界时下标为(r[i]),(r[i])为(i)二进制翻转后的数

我们可以先把系数分到底层的位置，然后一层一层往上做。

//Luogu P3803 多项式乘法 - 迭代FFT 
#include <complex>
#include <cstdio>
using namespace std;

typedef complex<double> comp;

const int N = (1 << 21) + 10;
const double PI = acos(-1);

int read() {
	int x = 0; char c = getchar();
	for(; c < '0' || c > '9'; c = getchar()) ;
	for(; c >= '0' && c <= '9'; c = getchar())
		x = x * 10 + (c & 15);
	return x;
}

int n, m, lim, r[N];
comp a[N], b[N];

void fft(comp * a, int type) {
	for(int i = 0; i < lim; i ++)
		if(i < r[i]) swap(a[i], a[r[i]]);
	for(int i = 1; i < lim; i <<= 1) {
		comp x(cos(PI / i), type * sin(PI / i));
		for(int j = 0; j < lim; j += (i << 1)) {
			comp y(1, 0);
			for(int k = 0; k < i; k ++, y *= x) {
				comp p = a[j + k], q = y * a[j + k + i];
				a[j + k] = p + q; a[j + k + i] = p - q;
			}
		}
	}
}

int main() {
	n = read(), m = read();
	for(int i = 0; i <= n; i ++) a[i] = read();
	for(int i = 0; i <= m; i ++) b[i] = read();
	
	int l = 0;
	for(lim = 1; lim <= n + m; lim <<= 1) ++ l;
	for(int i = 0; i < lim; i ++)
		r[i] = (r[i >> 1] >> 1) | ((i & 1) << (l - 1));
		
	fft(a, 1), fft(b, 1);
	for(int i = 0; i <= lim; i ++) a[i] *= b[i];
	fft(a, -1);
	
	for(int i = 0; i <= n + m; i ++)
		printf("%d ", (int)(0.5 + a[i].real() / lim));
	return 0;
}

upd：我们有更好的写法，预处理单位根

typedef double db;
typedef long long ll;
const db pi = acos(-1.0);
const int N = 1 << 21 | 5; // 2n^2
const int M = 1 << 22 | 5;
struct comp_t {
	db x, y;
	comp_t() {}
	comp_t(db a, db b) : x(a), y(b) {}
	comp_t operator + (const comp_t &b) const { return comp_t(x + b.x, y + b.y); }
	comp_t operator - (const comp_t &b) const { return comp_t(x - b.x, y - b.y); }
	comp_t operator * (const comp_t &b) const { return comp_t(x * b.x - y * b.y, b.x * y + x * b.y); }
	inline void rev() { y = -y; }
} wnk[M];
int lim, r[M], A[M], B[M], n;
ll ans[N];
void init(int len) {
	int l = 0;
	for(lim = 1; lim <= len; lim <<= 1) l ++;
	for(int i = 1; i < lim; i ++) {
		r[i] = (r[i >> 1] >> 1) | ((i & 1) << (l - 1));
	}
	wnk[0] = comp_t(1, 0);
	wnk[1] = comp_t(cos(2 * pi / lim), sin(2 * pi / lim));
	for(int i = 2; i < lim; i ++) wnk[i] = wnk[i - 1] * wnk[1];
}
void fft(comp_t *a, int t) {
	for(int i = 1; i < lim; i ++)
		if(i < r[i]) swap(a[i], a[r[i]]);
	for(int i = 1; i < lim; i <<= 1) {
		for(int j = 0; j < lim; j += i << 1) {
			for(int k = 0; k < i; k ++) {
				comp_t y = wnk[lim / (i << 1) * k];
				if(t == -1) y.rev();
				comp_t t = y * a[j + k + i];
				a[j + k + i] = a[j + k] - t;
				a[j + k] = a[j + k] + t;
			}
		}
	}
}
ll *const_mul(int *a, int *b, int n, int m) {
	static ll ans[M];
	static comp_t A[M], B[M];
	init(n + m);
	for(int i = 0; i < lim; i ++) A[i] = comp_t(i <= n ? a[i] : 0, 0);
	for(int i = 0; i < lim; i ++) B[i] = comp_t(i <= m ? b[i] : 0, 0);
	fft(A, 1); fft(B, 1);
	for(int i = 0; i < lim; i ++) A[i] = A[i] * B[i];
	fft(A, -1);
	for(int i = 0; i <= n + m; i ++) ans[i] = (ll) (A[i].x / lim + 0.5);
	return ans;
}

NTT 快速数论变换

（若不了解原根的基础知识，可以阅读 Hongzy：数学相关 中的内容）

对于一类特殊形式质数(p=r 2^k+1)，模(p)意义下做多项式乘法可以使用NTT。

回顾一下我们用到单位根三条性质：

1. (omega_n^n=omega_n^0=1)（逆变换中所用）
2. (omega_n^0,omega_n^1,omega_n^2,...,omega_n^{n-1})互不相同（点值表示的要求）
3. (omega_{2n}^{2k}=omega_{n}^{k})，(omega_{n}^{frac{n}{2}+k}=-omega_n^k)（分治基础）

根据(nmid p-1)启发，我们找到一个东西(g_n=g^{frac{p-1}{n}})（即(g^r)），也满足这三条性质：

1. (g_n^n=g_n^0=1)（证明：费马小定理）
2. (g_{n}^0,g_{n}^1,g_{n}^2,...,g_{n}^{n-1})互不相同（根据原根的性质，(g_n^0)到(g^{p-2})都是不同的）
3. (g_{2n}^{2k}=g^{frac{2k(p-1)}{2n}}=g^{frac{k(p-1)}{n}}=g_n^k,g_n^{frac{n}{2}+k}=g_n^k g_n^{frac{n}{2}}=-g_n^{k})

注意(3)的最后一步用到了(g_n^{frac{n}{2}}=g^{frac{p-1}{n}frac{n}{2}}=g^{frac{p-1}{2}})。由于(g)不是二次剩余，(g^{frac{p-1}{2}}=-1)

代码：

//UOJ#34. 多项式乘法
#include <algorithm>
#include <cstdio>
using namespace std;

const int MOD = 998244353;
const int N = 1e6 + 10;

int n, m, L, r[N], g[N], a[N], b[N];

int qpow(int a, int b) {
	int ans = 1;
	for(; b; b >>= 1, a = a * 1ll * a % MOD)
		if(b & 1) ans = ans * 1ll * a % MOD;
	return ans;
}

void ntt(int *a, int f) {
	for(int i = 0; i < n; i ++)
		if(i < r[i]) swap(a[i], a[r[i]]);
	for(int i = 1; i < n; i <<= 1) {
		int gn = qpow(3, (MOD - 1) / (i << 1)), x, y;
		for(int j = 0; j < n; j += (i << 1)) {
			int g = 1;
			for(int k = 0; k < i; k ++, g = 1ll * g * gn % MOD) {
				x = a[j + k]; y = 1ll * g * a[i + j + k] % MOD;
				a[j + k] = (x + y) % MOD;
				a[i + j + k]=(x - y + MOD) % MOD;
			}
		}
	}
	if(f == 1) return;
	reverse(a + 1, a + n);
	int y = qpow(n, MOD - 2);
	for(int i = 0; i < n; i ++) 
		a[i] = 1ll * a[i] * y % MOD;
}

int main() {
	scanf("%d%d", &n, &m);
	for(int i = 0; i <= n; i ++) 
		scanf("%d", &a[i]);
	for(int i = 0; i <= m; i ++) 
		scanf("%d", &b[i]);
	m += n;
	for(n = 1; n <= m; n <<= 1) L ++;
	for(int i = 0; i < n; i ++)
		r[i] = (r[i >> 1] >> 1) | ((i & 1) << (L - 1));
	ntt(a, 1); ntt(b, 1);
	for(int i = 0; i < n; i ++) 
		a[i] = 1ll * a[i] * b[i] % MOD;
	ntt(a, -1);
	for(int i = 0; i <= m; i ++)
		printf("%d ", a[i]);
	return 0;
}

结语

参考博客：

• Algocode 的算法之路
• RabbitHu's Blog
• Attack's Blog
• zzd's Blog