分治艺术的顶点，快速傅里叶变换

Chaofan

2022-03-04

| Views: | Total Words: 29k | Reading Time: 26 mins.

前言

快速傅里叶变换（Fast Fourier Transformation，简称 FFT），用于大整数乘法与多项式乘法/卷积的高效计算，其复杂度可以达到非常令人惊叹的 $O(n \log n)$ 。之所以这么说，就是因为这个问题看起来好像就感觉只能用 $O(n^2)$ 的算法来解决，并且直觉上似乎并没有什么优化空间的样子。而 FFT 用堪称艺术的巧妙取点和分治，非常优雅地解决了这个问题。

为什么突然想写这个算法呢？直接原因是，这一算法是我的算法课 “Divide and Conquer” 这一章所讲的最后一个算法。当然，我早已经不是第一次听到了——高中搞 OI 时候就学过一遍；在面试 ACM 班时被问及我最喜欢的算法时，我不假思索地回答了 FFT；上学期的科学计算课程（看起来和 FFT 八竿子打不着的数学课）也讲到了它，足以说明这一算法的魅力所在。可以说，我和它的缘分也是不浅了。

我一直认为，只要一个算法，或者说一个想法，具有令人惊叹、或者可以让人发自内心地赞叹“优美、巧妙”的地方，就已经有足够的价值了；至于它是否能被应用，其实也没那么重要。但 FFT 不仅巧妙绝伦，它还非常有用——因为它解决的问题实在是太基本了。

这大概就是这篇文章诞生的目的了。如果你此时能为“不知道 FFT”感到遗憾的话，那就看下去吧！

问题

因为大整数乘法其实本质就是多项式乘法，所以这里我们就直接来研究多项式乘法问题。

给定两个多项式 $F(x) = \sum_{i=0}^n a_i x^i $, $G(x) = \sum_{i=0}^m b_i x^i $, 求 $F(x) \cdot G(x)$. 你只要输出这个相乘后的多项式的系数就好。

Brute-Force

直观想法，当然是直接相乘啦~

我们考虑相乘后多项式项 $x^k$ 前的系数，可以采取一个比较组合的观点来看：可以是 $F(x)$ 中的 1 乘以 $G(x)$ 中的 $x^k$, 也可以是$F(x)$ 中的 $x$ 乘以 $G(x)$ 中的 $x^{k-1}$ … …

我们把这个想法写成式子，就是

$\sum_{i+j=k} a_i b_j x^k$

这个复杂度是多少呢？考虑相乘后的多项式长度为 $n+m-1$，我们枚举这个多项式的每一项，且每一项前面都有一个求和号，所有满足条件的 $i+j=k$ 有 $k+1$ 对，因此总复杂度为 $O(n^2)$（这里为了方便，我们不妨认为 $n$ 与 $m$ 是同一量级）

Karatsuba

这个优化方法是我在算法课上学到的，虽然不及 FFT 优秀，但是能将 n 上的系数从 2 优化成 1.x，而且想法也比较简单。

考虑我们计算两位数乘法

$\overline{ab} \times \overline{cd} = (10a+b)(10c+d) = 100ac+10ad+10bc+bd = 100ac+10(ad+bc)+bd$

那么两位数乘法需要花费 4 单位的时间。我们考虑式子展开后，需要计算 $ac, ad, bc, bd$ 四个一位数乘法，也需要 4 单位的时间。

算法的关键部分来了：我们其实并不需要计算 4 个东西的乘法，因为我们实际关心的是 $ac, ad+bc, bd$ 三个式子。我们考虑如果我们计算 $(a+b)(c+d), ac, bd$ 三个乘法，是不是就可以通过加减法得到我们要的三个式子了。

当然直接用两位数乘法来讲，可能会导致你产生一些关于计算模型上的疑惑（比如乘以 100 要花费多少时间？$(a+b)(c+d)$ 这一项是一位数乘法还是两位数？）这些都不重要，我们直接把这个思想用于多项式中，这样你就知道这一算法的优化之处了。

为了方便暂时设两个多项式长度相同，记

$F(x) = f_1(x) + f_2(x) x^{n/2},\ G(x) = g_1(x) + g_2(x) x^{n/2}$

那么

$F(x) \cdot G(x) = [f_1(x) + f_2(x) x^{\frac{n}{2}}] [g_1(x) + g_2(x) x^{\frac{n}{2}}] \\ = f_1(x) g_1(x) + [f_2(x) g_1(x) + f_1(x) g_2(x)] x^{\frac{n}{2}} + f_2(x) g_2(x) x^n$

其中 $f_1, f_2, g_1, g_2$ 均为长度为 $\frac{n}{2}$ 的多项式。应用上面的方法，我们也只要计算三个多项式的乘积：$(f_1+f_2)(g_1+g_2), f_1g_1, f_2g_2$. 这样我们的时间复杂度递推式就可以写作

$T(n) = 3T(\frac{n}{2}) + O(n)$

应用 Master Theorem（当然你也可以用你高超的数学技巧）可以求出复杂度为 $O(n^{\log_2 3})$，大约是 $O(n^{1.6})$. 这一算法被称为 Karatsuba 算法。

为什么要点值表示

前面讲了那么多其它做法，始终没有来到正题。

虽然 Karatsuba 已经比较巧妙了，但我们仍能感觉到其中的瓶颈——你总是得算点什么东西。并且已经证明了每次算“3”个多项式已经是最少的了（如果你能做到每次转化成算两个 $T(\frac{n}{2})$ 的操作，那你就做到了线性复杂度！）

所以我们需要来一个思路上的转变。FFT 第一个能够优化的关键，就是它使用点值表达，而非系数表达的方式来考虑一个多项式，进而能够使得乘法操作变得非常简单。

什么叫点值表达呢？我们知道一个多项式的所有系数，当然可以确定它；但我们只要知道其之上足够多的点，我们也能确定它。事实上，对于 n 次多项式，我们只要知道 n+1 个不同的点就可以做到。下面我们进行一个简单的证明。

考虑现在你有点值 $(x_0, F(x_0)), (x_1, F(x_1)), \dots, (x_n, F(x_n))$，我们先做一个待定系数

$F(x) = a_0 + a_1 x + \dots + a_n x^n$

我们把系数当成未知数，只要我们的信息量足够解出它们就可以了。因此我们得到方程

$\left[\begin{matrix} F(x_0) \\ F(x_1) \\ \vdots \\ F(x_n) \end{matrix}\right] = \left(\begin{matrix} 1 & x_0 & \dots & x_0^n \\ 1 & x_1 & \dots & x_1^n \\ \vdots & \vdots & & \vdots \\ 1 & x_n & \dots & x_n^n \end{matrix}\right) \left[\begin{matrix} a_0 \\ a_1 \\ \vdots \\ a_n \end{matrix}\right]$

注意到这是一个 Vandermonde 矩阵，因此其可逆。所以此方程有唯一解，这就证明了 n+1 个点的点值信息可以唯一确定此多项式。

但其实我们日常生活中谈论“多项式”，点值表示其实非常少用。毕竟在做数学作业时，你也不会拿着一些点，然后称其为“多项式”吧？所以在这里特地提到这一表示法，它肯定有其优势之处——在解决我们的乘法问题方面。

考虑一开始的乘法问题 $F(x) \cdot G(x)$，我们现在如果用点值的观点去思考，就不用去算它的系数了——我们只要找到相当于它“次数+1”（在我们的预设中，即为 $n+m+1$）个点。那我们只要选定 $x_0, x_1, \dots, x_{n+m+1}$ 个横坐标，求出 $F(x_i)$ 与 $G(x_i)$ 的点值，那么乘法结果在此处的点值即为 $F(x_i) \cdot G(x_i)$.

你可以注意到，我们的乘法 $F(x_i) \cdot G(x_i)$ 在这时只需要 $n+m+1$ 次！我们采用点值表示的方法，绕开了多项式系数那些“相互组合”的情况，比较本质地把这个“乘法”给体现出来了。

但坏消息是，得到 $F(x_i)$ 与 $G(x_i)$ 似乎不是个很高效的事情。简单来想，我们将一个值代入 n 次多项式，似乎时间复杂度不可能低于 $O(n)$。所以这样想来，好像求出所有的 $F(x_i)$ 与 $G(x_i)$ 复杂度也是 $O(n^2)$。不过，这似乎看起来要比简单的系数乘法要有前景一点。而优化这一步，正是 FFT 的最核心之处。

DFT-高效地从系数表示到点值表示

我们在回顾上一步，现在我们第一步需要从系数表示快速地求得点值表示。

重新考虑这个方程（实际上，这个方程就是沟通两个表示法的方程）

写成这样，也就不难理解我们上一部分所说的 $O(n^2)$ 复杂度了——暴力地进行矩阵乘以向量，当然是平方啦。但 FFT 是怎么优化到 $O(n \log n)$ 的呢？

设计算法第一原则：每当我们想要优化算法时，想想我们现在的算法重复干了什么事，以及有什么信息没有用到。（当然，这是我自己编的）

似乎看上去，我们并没有在重复干事情。但我们可以很快意识到，我们忽略了一些条件，或者说变相加强了问题——我们在研究任意取点 $x_0, x_1, \dots, x_n$ 的情况，但取点完全可以由我们自己决定。

首先，我们肯定会趋向代入一些，尽量让计算简单的情况——比如 $x=0$。算都不用算！但你需要继续找 $n+m$ 个像 0 这样特殊的点，这对于数据规模越来越大的情况实在不是很可行。

所以，我们不能想着找“若干个”特殊点，而应该找“若干类”，或者说想办法让点与点之间有关联、有大量重复部分，进而简化计算。

考虑我们取 $x_0 = 1, x_1 = -1$ 两个（看起来就很有关系）的对偶点，代入得到

$F(1) = a_0 + a_1 + \dots + a_n \\ F(-1) = a_0 - a_1 + \dots + (-1)^n a_n$

假如我们记奇数部分的和为 $f_1$，偶数部分的和为 $f_2$，那么有 $F(1)=f_1+f_2, F(-1)=f_1-f_2$。而计算 $f_1, f_2$ 的总和时间为 $n$（而不是 $2n$！），也就意味着这两个点值的计算，我们只花费了一个 $n$。

而这样的点对关系其实并不特殊——任意两个互为相反数的取点都能做到。假设我们能够对于 $f_1$ 与 $f_2$ 这种部分不断往下如此做，即

$F(x) = (a_0+a_2 x^2 + a_4 x^4 + \dots) + x (a_1 + a_3 x^2 + a_5 x^4 + \dots) \\ F(-x) = (a_0+a_2 x^2 + a_4 x^4 + \dots) - x (a_1 + a_3 x^2 + a_5 x^4 + \dots)$

我们就能得到递推式

$T(n) = 2 T(\frac{n}{2}) + O(n)$

注意这里的 n 表示的是点的个数（按上面的假设其实应该是 n+1，不过讨论复杂度并不用关心），我们本来要求 n 个点的点值，但因为我们的取点是 $x_0, -x_0, x_1, -x_1, \dots$，我们只需要计算其中一半（比如都为正的那一半）的取点情况。

注意上面有个前提：“如果我们能一直做下去”。假设这个成立的话，根据上述递推式也很容易求出时间复杂度为 $O(n \log n)$。但很遗憾，这个前提明显是不成立的——事实上，应该是“一定不”。因为假设原问题的点是 $x_0, -x_0, x_1, -x_1, \dots$，那么子问题的点就是 $x_0^2, x_1^2, \dots$，如果子问题还能像这样“相反数配对”的话，即有 $x_i^2=-x_j^2$，得到 $x_i=x_j=0$（如果是在实数域讨论的话）

相信读者已经发现端倪了——实数域确实无法找到，但在复数域似乎大有可为。如果考虑复数域，$x_i^2=-x_j^2$ 即为 $x_i=i x_j$（注意 i 在此处有点重名，注意辨别复数单位 i 与下标 i）

因此假设我们需要快速对 4 个取点的情况进行计算，首先我们当然取 $1, -1$ 两个比较方便的相反数。又根据上面，$i, -i$ 也要纳入我们的考量。我们进行一次子问题转换：$1, -1, i, -i$ 变为 $1, -1$，能够继续进行下去！

事实上，这些特殊的点就是复数域上的 $n$ 次单位根。“相反数”在复数的空间上其实就是相差 $ \pi $ 的角度，而“乘以 i”则相当于相差 $ \frac{\pi}{2} $ 的角度。所以一组符合上面条件的 4 个取点实际上就是一个“十字架”。

我们可以再来手模一下 8 个点的过程，我们取 8 次单位根，按照正负-正负分组，每次进行一次平方操作往下变换：

$\begin{matrix} 1 & -1 & i & -i & \sqrt{i} & -\sqrt{i} & \sqrt{-i} & -\sqrt{-i} \\ 1 & -1 & i & -i \\ 1 & -1 \\ 1 \end{matrix}$

这样，我们就做到了在 $O(n \log n)$ 中高效地转换点值表达式了。这一过程被称为 DFT（Discrete Fourier Transform，离散傅里叶变换）。

接下来，相乘的时间复杂度是线性的，所以我们似乎已经得到了结果——但是之前说过，实际生活中通常期望我们的结果表示成系数的表达方法，因此我们还需要考虑能否高效地将点值表达式转回系数表达式。

IDFT-高效地从点值表示到系数表示

重新考虑那个方程，为了方便我们记函数值的向量为 $Y$, 矩阵为 $X$, 系数向量为 $\alpha$，将方程重写为

$Y = X \alpha$

我们知道取 n 次单位根作为横坐标时（记此时的矩阵 $X$ 为一个特殊的矩阵 $W$），我们已经能高效从 $\alpha$ 得到 $Y$。而现在相当于我有 $Y$，并且矩阵还是 $W$（这里不能再自己取值了，因为我们求乘积所用的 $F$ 和 $G$ 的点值表达所取的横坐标是 n 次单位根，那乘积的点值表达也只能是这些横坐标）

已经证明 $W$ 是可逆的，因此我们只需要能够比较快速地计算

$\alpha = W^{-1} Y$

这里涉及到两个难题：求逆和乘法。我们知道矩阵求逆可不是一件轻松的事，一般的矩阵需要 $O(n^3)$。但我们既然都跟着 FFT 做到这一步了，说明 $W$ 一定有某种奇妙的性质使得它可以快速求逆。

这里就要点一下题了，我们来看 FFT 的全称“快速傅里叶变换”。学过分析的同学可能都知道傅里叶变换，之所以 FFT 和分析里的傅里叶变换本质是一样的，是因为复数域的单位根和三角函数在某种意义上是一样的——其实分析里的傅里叶变换用复数的形式写，就很能看出它们的关系了。

因此，我们不难从傅里叶变换来猜想，$W$ 是否可能具有某种正交性？如果其为正交矩阵，那它求逆可太简单了——就是它的共轭转置（正交矩阵满足 $W^H W = E$）。

观察到 $W$ 其实是个对称矩阵，因此 $W^H$ 只要把其元素取共轭就好了。我们观察 $W$ 的形态（假设对于取 $n+1$ 个点的情况，我们取 $n+1$ 次单位根）

$\left(\begin{matrix} 1 & 1 & 1 & \dots & 1 \\ 1 & \omega & \omega^2 & \dots & \omega^n \\ 1 & \omega^2 & \omega^4 & \dots & \omega^{2n} \\ \vdots & \vdots & \vdots & & \vdots \\ 1 & \omega^n & \omega^{2n} & \dots & \omega^{n^2} \end{matrix}\right)$

因此 $W^HW$ 的 $(i, i)$（也即对角线位置一定是）

$1 + \overline{\omega^{i+1}}\omega^{i+1} + \overline{\omega^{2(i+1)}}\omega^{2(i+1)} + \dots + \overline{\omega^{n(i+1)}}\omega^{n(i+1)} = n+1$

而非对角线部分，$(i, j)$

$\begin{aligned} &1 + \overline{\omega^{i+1}}\omega^{j+1} + \overline{\omega^{2(i+1)}}\omega^{2(j+1)} + \dots + \overline{\omega^{n(i+1)}}\omega^{n(j+1)} \\ =\ &1 + \omega^{i-j} + \omega^{2(i-j)} + \dots + \omega^{n(j-i)} \\ =\ &\frac{1-\omega^{(n+1)(j-i)}}{1-\omega^{j-i}} = 0 \end{aligned}$

这里你需要几个基本的数学知识

对于一个复数 $\omega$，$ \overline{\omega} \omega = \omega ^2 $。上文中均为单位根，所以模长为 1。
对于复数我们也可以应用等比数列求和公式。
对于 $n+1$ 次单位根，$ \omega^{n+1} = 1 $。

因此我们得到其“基本是正交的”（在分析里，其实可以认为这就是正交的；不过在矩阵的定义中，正交矩阵相乘一定要等于 $E$ 而非 $cE$）。也即我们已经证明了

$W^H W = (n+1) E$

所以我们有

$\alpha = W^{-1} Y = \frac{1}{n+1} W^H Y$

所以求逆我们实际上不需要花时间。但是我们还需要考虑矩阵乘法的事。我们来考量 $W^H$ 的形态：

$\left(\begin{matrix} 1 & 1 & 1 & \dots & 1 \\ 1 & \omega^{-1} & \omega^{-2} & \dots & \omega^{-n} \\ 1 & \omega^{-2} & \omega^{-4} & \dots & \omega^{-2n} \\ \vdots & \vdots & \vdots & & \vdots \\ 1 & \omega^{-n} & \omega^{-2n} & \dots & \omega^{-n^2} \end{matrix}\right)$

其中共轭用单位根的负指数表示，因为共轭实际上就是等价于某种倒数，或者说乘法逆元。我们惊喜地发现这和 $W$ 相比，只是把所有指数都换成了负号而已！那我们已经可以高效计算 $W$ 与一个向量相乘了，那$W^H$ 与一个向量相乘，我们似乎也应该能够很高效地计算——事实上确实是这样的。

这是因为，我们能让算法高效进行的关键部分都保留住了。我们能让这个相乘高效进行，本质就是因为单位根良好的旋转对称性，而取负的单位根无非就是旋转方向从逆时针变成了顺时针，其它没有任何改变。所以如果你的 DFT 函数接口设计得好的话，你只要调用你已经写好的 DFT，只要在传参时候把单位根的选取变成它们的共轭，就可以完成这个过程。

这个过程被称为 IDFT，也就是 DFT 的逆变换。至此，我们的 FFT 算法就完成了。

思考总结

讲了这么多，还是觉得很不可思议。到底是什么性质能让我们如此高效地计算呢？我们可以做一下简单的归纳总结。首先，我们可以先梳理一下 FFT 的算法过程。

Step 1. DFT，用 n 次单位根高效计算 F 与 G 的 n+m+1 个点值。
Step 2. 乘法。将 F 与 G 的对应点值相乘，得到 $F \cdot G$ 的 n+m+1 个点值。
Step 3. IDFT。将 $F \cdot G$ 还原成系数形式。

我们发现，点值表达式的优越性是这个算法的关键一步——这是一个很棒的想法。它启示我们，其实当我们要求一个东西时，并不一定要按定义将其求出来，我们只要致力于得到“能够确定目标的信息量”就足够了。

当然，FFT 算法的灵魂还是 n 次单位根的使用。其能进行高效地分治，本质原因是良好的对称性；而在 IDFT 化归为 DFT 的分析中，我们又用到了它的正交性。这才是 FFT 真正实现 $n \log n$ 最为重要的一环。

NTT

虽然 FFT 已经很完美了，但是在实际实现过程中，由于复数域的 n 次单位根通常需要用实数表示，所以存在精度误差的问题。

NTT（Number Theory Transformation，数论变换）则使用数论中的“原根”来代替 FFT 中的单位根。在模意义下，原根具有和单位根相同的性质，并且因为 NTT 全程均使用整数，消除了精度误差，因此也常用于替代 FFT。

代码

可以通过洛谷模板题的代码。

#include <cstdio>
#include <vector>
#include <cmath>

using namespace std;

const double PI = acos(-1.0);
const int N = (1<<22);

struct Complex {
    double re, im;

    Complex() : re(0.0), im(0.0) {}

    Complex(double re_, double im_) : re(re_), im(im_) {}

    friend Complex operator + (Complex obj1, Complex obj2) {
        return Complex(obj1.re + obj2.re, obj1.im + obj2.im);
    }

    friend Complex operator - (Complex obj1, Complex obj2) {
        return Complex(obj1.re - obj2.re, obj1.im - obj2.im);
    }

    friend Complex operator * (Complex obj1, Complex obj2) {
        return Complex(obj1.re * obj2.re - obj1.im * obj2.im, obj1.re * obj2.im + obj1.im * obj2.re);
    }
} f[N], g[N], result[N]; // 使用复数存储系数, 统一接口

// 得到第一个 >= num 的 2^n
// 因为要把 FFT 的工作长度对齐到 2 的幂次
void get2Power(int& result, int num) {
    result = 1;
    while (result < num) {
        result <<= 1;
    }
}

// 这个函数相当于一个DFT, 不过IDFT也要用, 所以直接叫做FFT
// 分治地计算 W * vector, 时间 O(n log n)
// arr 一开始存储系数, 计算完点值后直接覆盖 arr 返回 (反正系数没用了)
// isIDFT 表示在 IDFT, 如果是的话要用单位根的共轭做
void FFT(Complex* arr, int len, bool isIDFT) {
    if (len <= 1) { // 如果 f(x) = 3, 那么显然 f(x0) = 3 就是点值表达
        return ;
    }

    // 奇偶系数拷贝
    Complex *odd = new Complex [len>>1];
    Complex *even = new Complex [len>>1];

    for (int i = 0; i < len; i++) {
        if (i & 1) odd[i >> 1] = arr[i];
        else even[i >> 1] = arr[i];
    }

    FFT(odd, len>>1, isIDFT);
    FFT(even, len>>1, isIDFT);

    // 构造 len 次单位根 w
    Complex w = Complex(cos(2*PI/len), sin(2*PI/len));
    // 储存 w 的 n 次幂
    Complex wPower = Complex(1, 0); 
    if (isIDFT) w.im = -w.im; // 取共轭

    // f(x) = a0 + a1 x + a2 x^2 + ...
    //      = (a0 + a2 x^2 + ...) + x (a1 + a3 x^2 + ...)
    //      = f1(x) + x f2(x)
    // f(-x) = (a0 + a2 x^2 + ...) - x (a1 + a3 x^2 + ...)
    //       = f1(x) - x f2(x)
    // 代入的点依次是 1, w, w^2, ..., 所以 wPower 就是当前带入的 x
    // f1 是 even, f2 是 odd
    // 注意这里合并时的顺序, 我们采用 1 i -1 -i (旋转顺序) 而不是 i -1 i -i

    for (int i = 0; i < (len >> 1); i++) {
        arr[i] = even[i] + wPower * odd[i];
        arr[i+(len>>1)] = even[i] - wPower * odd[i];
        wPower = wPower * w;
    }

    printf("fft result, len = %d\n", len);
    for (int i = 0; i < len; i++) 
        printf("result re = %lf, im = %lf\n", arr[i].re, arr[i].im);
    printf("\n");

    delete [] odd;
    delete [] even;
}

int n, m, resultLen;

int main() {
    // printf("N = %d\n", N);

    // test complex
    // Complex a = Complex(3, 4);
    // Complex b = Complex(6, 5);
    // Complex c = a * b;
    // printf("c = %lf + i %lf", c.re, c.im);

    // test PI
    // printf("pi = %lf", PI);

    scanf("%d %d", &n, &m);
    for (int i = 0; i <= n; ++i) scanf("%lf", &f[i].re);
    for (int i = 0; i <= m; ++i) scanf("%lf", &g[i].re);

    get2Power(resultLen, n+m+1);

    // printf("resultLen = %d\n", resultLen);

    // 对齐之后超过的部分都是 0 不影响计算
    FFT(f, resultLen, false);
    FFT(g, resultLen, false);

    // 点值表达式的计算
    for (int i = 0; i < resultLen; ++i)
        result[i] = f[i] * g[i];

    // printf("result = \n");
    // for (int i = 0; i < resultLen; i++)
    //     printf("i=%d, re=%lf, im=%lf\n", i, result[i].re, result[i].im);

    // IDFT
    FFT(result, resultLen, true);

    // 题目要求输出 n+m 个数
    // 注意正交那一步要除掉一个常数, 因为神秘精度问题还要 +0.5
    // result 理应是实数
    for (int i = 0; i <= n+m; ++i)
        printf("%d ", int(result[i].re/resultLen + 0.5));

    return 0;
}

Author: SiriusNEO

Published on: Metric Space

All posts on this blog are licensed under the CC BY-NC-SA 4.0 license unless otherwise noted.