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

前言

快速傅里叶变换（\(\text{Fast Fourier Transform,FFT}\) ）是一种能在\(O(n \log n)\)的时间内完成多项式乘法的算法，在\(OI\)中的应用很多，是多项式相关内容的基础。下面从头开始介绍\(\text{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}\).实轴到复数向量的转角\(\theta\)称为幅角。

复数的基本运算

复数的加（减）法：\((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\)倍是\(2\pi\)的倍数

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

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

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

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

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

单位根均分复平面单位圆，那\(w_n\)的幅角就是\(\frac{2\pi}{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{2\pi 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)\)的离散傅里叶变换（\(\text{Discrete Fourier Transform, DFT}\)），写作\(\vec y=\text{DFT}_n(\vec a)\)。

直接求\(\text{DFT}\)是\(O(n^2)\)的。\(\text{FFT}\)的常用算法\(\text{Cooley-Tukey}\)使用分治方法做到\(O(n\log 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_i\omega_{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^k\sum_{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

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

结论：只要把\(\text{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)\not = 1\)即\(k\not = 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\not = 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。

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

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

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

\(g_n^n=g_n^0=1\)（证明：费马小定理）
\(g_{n}^0,g_{n}^1,g_{n}^2,...,g_{n}^{n-1}\)互不相同（根据原根的性质，\(g_n^0\)到\(g^{p-2}\)都是不同的）
\(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;

}

结语

参考博客：