快速傅里叶变换(FFT)算法的实现及优化

By Long Luo

之前的文章 快速傅里叶变换(FFT)算法 详细介绍了 FFT\textit{FFT} 的具体实现,也实现了分治版FFT\textit{FFT}

分治版 FFT\textit{FFT} 代码很容易理解,但递归会消耗 O(logn)O(logn) 的栈空间,同时代码实现中还有很多优化空间,这篇文章我们就来优化 FFT\textit{FFT} 下的实现。

Cooley–Tukey FFT Algorithm

Cooley–Tukey FFT Algorithm\textit{Cooley–Tukey FFT Algorithm} [1]通过使用分治算法[2]实现将复杂问题简单化。

FFT 分治过程

FFT

具体操作流程可以分为 33 步:

  1. Divide:按照奇偶分为 22 个子问题。
  2. Conquer:递归处理 22 个子问题
  3. Combine:合并 22 个子问题。

蝴蝶操作

分治的前 22 步之前已经详细讲述过了,第 33 步合并操作是通过**蝴蝶操作(Butterfly Operations)**来实现的,其示意图如下所示:

Butterfly Operations

蝴蝶操作可能会让人看了一头雾水,不过没关系,我们一步一步来,彻底弄懂她!

空间优化

回顾之前的 递归版代码 ,在合并这一步,这一层有 nn 项需要处理,我们新建了一个数组 y(n)y(n),这是为什么呢?

我们可以复用之前的 yey_eyoy_o 数组以降低空间复杂度吗?

先看下合并需要做的操作:

y(k)=ye(k)+ωnkyo(k+n2)y(k) = y_e(k) + \omega_{n}^{k} \cdot y_o(k+\frac{n}{2})

y(k+n2)=ye(k)ωnkyo(k+n2)y(k + \frac{n}{2}) = y_e(k) - \omega_{n}^{k} \cdot y_o(k+\frac{n}{2})

很明显,我们如果复用 yey_eyoy_o 数组的话,那 y(k)y(k)y(k+n2)y(k + \frac{n}{2}) 至少有一个数据会受影响,所以我们需要额外的 y(n)y(n) 数组来存储数据。

那么有没有办法来做到复用 yey_eyoy_o 数组呢?

当然可以!

我们只要将修改下合并顺序,加入一个临时变量 t=ωnkye(k+n2)t=\omega_{n}^{k} \cdot y_e(k+\frac{n}{2}) ,合并过程就可以在原数组中进行:

cd  t=ωnkye(k+n2)cd \; t = \omega_{n}^{k} \cdot y_e(k+\frac{n}{2})

ye(k+n2)=ye(k)ty_e(k+\frac{n}{2}) = y_e(k) - t

ye(k)=ye(k)+ty_e(k) = y_e(k)+t

这样就可以原地进行了,不再需要额外数组。

位逆序置换

在分治过程中,每一次都会把整个多项式的奇数次项和偶数次项系数分开,一直分到只剩下一个系数。因此,如果我们可以先模拟递归把这些系数在原数组中交换位置,把每个系数都放在最终的位置上,再一步一步向上合并。

88 项多项式为例,模拟递归拆分的过程:每一次分组情况如下图:

  • 初始序列为 {x0,x1,x2,x3,x4,x5,x6,x7}\{x_0, x_1, x_2, x_3, x_4, x_5, x_6, x_7\}

  • 一次二分之后 {x0,x2,x4,x6},{x1,x3,x5,x7}\{x_0, x_2, x_4, x_6\}, \{x_1, x_3, x_5, x_7\}

  • 两次二分之后 {x0,x4},{x2,x6},{x1,x5},{x3,x7}\{x_0, x_4\}, \{x_2, x_6\}, \{x_1, x_5\}, \{x_3, x_7\}

  • 三次二分之后 {x0}{x4}{x2}{x6}{x1}{x5}{x3}{x7}\{x_0\} \{x_4\} \{x_2\} \{x_6\} \{x_1\} \{x_5\} \{x_3\} \{x_7\}

拆分前后数组位置变化情况如下:

Index
原序列: 0 1 2 3 4 5 6 7
原数组: 000 001 010 011 100 101 110 111
后序列: 0 4 2 6 1 5 3 7
后数组: 000 100 010 110 001 101 011 111

发现规律了没?

二进制翻转

后序列是原序列的二进制翻转。原序列每个数用二进制表示,然后把二进制翻转对称一下,就是最终那个位置的下标。比如 x1=(001)2x_1=(001)_{2} ,翻转是 4=(100)24=(100)_2 。这个变换为位逆序置换(bit-reversal permutation,也就是蝴蝶变换)。

根据它的定义,我们可以在 O(nlogn)O(nlog n) 的时间内求出每个数变换后的结果:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
/*
* 进行 FFT 和 IFFT 前的反置变换
* 位置 i 和 i 的二进制反转后的位置互换
* len 必须为 2 的幂
*/
void change(vector<complex<double>>& y, int len) {
int k;
for (int i = 1, j = len / 2; i < len - 1; i++) {
if (i < j) {
swap(y[i], y[j]);
}

// 交换互为小标反转的元素,i<j 保证交换一次
// i 做正常的 + 1,j 做反转类型的 + 1,始终保持 i 和 j 是反转的
k = len / 2;
while (j >= k) {
j = j - k;
k = k / 2;
}

if (j < k) {
j += k;
}
}
}

反向二进制加法

实际上,位逆序变换相当于一个反向二进制加法,因为原序列是每次从二进制右端加 11 地递增,翻转后,就相当于每次从二进制左端加 11,这样的话就可以 O(n)O(n) 从小到大递推实现。

我们来推理下如何实现[3]

len=2klen=2^k,其中 kk 表示二进制数的长度,设 R(x)R(x) 表示长度为 kk 的二进制数 xx 翻转后的数(高位补 00)。我们要求的是: R(0),R(1),,R(n1)R(0),R(1),\cdots,R(n-1)

  1. R(0)=0R(0)=0

  2. 从小到大求 R(x)R(x)

待求的问题可以分为 22 步:

  1. 求最后一位之前的二进制数的倒置结果;

  2. 把被抹去的原来的最后一位放回结果的第一位。

对于第一步:

最后一位前面的二进制数取反后是多少?因为我们进行位逆序置换的过程是从 00lenlen 的,而求最后一位前面的二进制数也就等于把所需要求的数的右移一位(除以 22),所以最后一位前面的二进制数的倒置就是你在之前已经求过的 R(x>>1)=R(x2)R(x >> 1) = R\left(\left\lfloor \dfrac{x}{2} \right\rfloor\right)

但我们在进行位逆序置换时,会将一个数的二进制表示前面的 00 补齐,使得其长度等于 kk,例如 k=5k=5(12)10=(1100)2(12)_10=(1100)_2 补齐至 (01100)2(01100)_2,再进行位逆序置换得到 (00110)2(00110)_2。这样我们得到的 R(x/2)R(x/2) 其实与我们所求的最后一位前面的二进制数的倒置是不一致的,我们要求 (1100)2(1100)_2 的倒置,但求得的却是 (01100)2(01100)_2 的倒置,也就是说我们在 xx 的最前面多了一位,逆置之后也就是在结果的最后多了一位,所以我们第一个问题的结果不是 R(x>>1)R(x >> 1) 而是 R(x>>1)>>1R(x >> 1) >> 1, 也就是通过将原来的结果右移一位消去补上去的那一位。

解决了第一步,对于第二步来说就非常简单了。如果之前的个位是 00,翻转之后最高位就是 00;如果个位是 11,则翻转后最高位是 11,也就是结果再加上 len2=2k1\dfrac{len}{2}=2^{k-1}

综上:

R(x)=R(x2)2+(xmod2)×len2R(x)=\left\lfloor \frac{R\left(\left\lfloor \frac{x}{2} \right\rfloor\right)}{2} \right\rfloor + (x\bmod 2)\times \frac{len}{2}

举个例子来说明下:

k=5k=5len=(32)10=(100000)2len=(32)_{10}=(100000)_2 ,假设待翻转的数是 (25)10=(11001)2(25)_{10}=(11001)_2

  1. 考虑 (12)10=(1100)2(12)_{10}=(1100)_2 ,我们知道 R((1100)2)=R((01100)2)=(00110)2R((1100)_2)=R((01100)_2)=(00110)_2 ,再右移一位就得到了 (3)10=(00011)2(3)_{10}=(00011)_2

  2. 考虑个位,如果是 11 ,它就要翻转到数的最高位,即翻转数加上 (16)10=(10000)2=2k1(16)_{10}=(10000)_2=2^{k-1} ,如果是 00 则不用更改;

  3. 最终结果是 (00011)2+(10000)2=(19)10(00011)_2 + (10000)_2 = (19)_{10}

位逆序变换代码实现:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
// 同样需要保证 len 是 2 的幂
// 记 rev[i] 为 i 翻转后的值
void bitReverse(vector<complex<double>> &y, int len) {
vector<int> rev(len);

rev[0] = 0;
for (int i = 1; i < len; i++) {
// 最后一位之前的部分
rev[i] = rev[i >> 1] >> 1;
if (i & 1) { // 如果最后一位是 1,需要再加上 len/2
rev[i] |= len >> 1;
}
}

for (int i = 0; i < len; ++i) {
if (i < rev[i]) { // 保证每对数只翻转一次
swap(y[i], y[rev[i]]);
}
}

return;
}

FFT 迭代版

通过位逆序置换,我们可以用迭代替代分治递归 FFT\textit{FFT} ,优化 FFT\textit{FFT} 性能。

回放原先的递归过程,方便理解(从下往上倒推):

{a0,a1,a2,a3,a4,a5,a6,a7}\left \{ a_0,a_1,a_2,a_3,a_4,a_5,a_6,a_7 \right \}

{a0,a2,a4,a6},{a1,a3,a5,a7}\left \{ a_0,a_2,a_4,a_6 \right \}, \left\{ a_1,a_3,a_5,a_7 \right \}

{a0,a4},{a2,a6},{a1,a5},{a3,a7}\left \{a_0,a_4 \right \},\left \{ a_2,a_6 \right \},\left\{ a_1,a_5 \right \},\left\{ a_3,a_7 \right \}

{a0},{a4},{a2},{a6},{a1},{a5},{a3},{a7}\left\{ a_0 \right \},\left\{ a_4 \right \},\left\{ a_2 \right \},\left\{ a_6 \right \},\left\{ a_1 \right \},\left\{ a_5 \right \},\left\{ a_3 \right \},\left\{ a_7 \right \}

这其中有三层循环,第一层循环表示当前需要合并的多项式的长度:从每段长度为 11 合并为 22,而后 44 直到最大:

{a0,a4}len=2,{a2,a6}len=2,{a1,a5}len=2,{a3,a7}len=2\underbrace{\left\{a_0,a_4 \right \}}_{len=2},\underbrace{\left\{ a_2,a_6 \right \}}_{len=2},\underbrace{\left\{a_1,a_5 \right \}}_{len=2},\underbrace{\left\{a_3,a_7 \right \} }_{len=2}

len<<1len<<1

{a0,a2,a4,a6}len=4,{a1,a3,a5,a7}len=4\underbrace{\left\{ a_0,a_2,a_4,a_6 \right \}}_{len=4}, \underbrace{\left\{ a_1,a_3,a_5,a_7 \right \}}_{len=4}

以此类推:

1
2
3
for (int len = 2; len <= n; len <<= 1) {

}

之后我们需要标记每个多项式的开头( j+=lenj+=len ):

j=0{a0,a4}len=2,j+=len{a2,a6}len=2,j+=len{a1,a5}len=2,j+=len{a3,a7}len=2\mathop{\uparrow}\limits_{j=0} \underbrace{ \left\{a_0,a_4 \right \} }_{len=2},\mathop{\uparrow}\limits_{j+=len}\underbrace{\left\{a_2,a_6 \right \}}_{len=2},\mathop{\uparrow}\limits_{j+=len}\underbrace{\left\{ a_1,a_5 \right \}}_{len=2},\mathop{\uparrow}\limits_{j+=len}\underbrace{\left\{ a_3,a_7 \right \}}_{len=2}

同时因为每个多项式对应的根不一样, 因此我们需要在两层循环之间定义单位根 ωn\omega_nω\omega,不断乘以 ωn\omega_n 切换根的数值。

1
2
3
4
5
6
7
8
9
double ang = 2 * PI / n * (invert ? -1 : 1);
complex<double> wn(cos(ang), sin(ang)); // omega为第一个n次复根,
complex<double> w(1, 0); // curr为第零0个n次复根, 即为 1

for (int i = 2; i <= n; i <<= 1) {
for (int j = 0; j < n; j += i) {

}
}

那么第三重循环的实现就很简单了:

1
2
3
4
5
6
7
8
9
10
11
for (int i = 2; i <= n; i <<= 1) {
for (int j = 0; j < n; j += i) {
for (int k = j; k < j + i / 2; k++) {
complex<double> t = a[k];
complex<double> b = w * a[k + i / 2];
a[k] = t + b;
a[k + i / 2] = t - b;
w *= wn;
}
}
}

FFT\textit{FFT} 迭代版代码最终实现如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
/**
* FFT Iteration 实现
*
* @param a
* @param invert true means IFFT, else FFT
* @return y
*/
void FFT(vector<complex<double>> &a, bool invert) {
int n = a.size();

if (n == 1) {
return;
}

bitReverse(a);

for (int len = 2; len <= n; len <<= 1) {
double ang = 2 * PI / len * (invert ? -1 : 1);
complex<double> wlen(cos(ang), sin(ang));

for (int i = 0; i < n; i += len) {
complex<double> w(1, 0);
for (int j = 0; j < len / 2; j++) {
complex<double> u = a[i + j];
complex<double> v = w * a[i + j + len / 2];
a[i + j] = u + v;
a[i + j + len / 2] = u - v;
w *= wlen;
}
}
}

if (invert) {
for (auto &x : a) {
x /= n;
}
}
}

应用及实践

43. 字符串相乘

POLYMUL - Polynomial Multiplication[4]

参考资料


  1. Wiki: Cooley–Tukey FFT algorithm ↩︎

  2. Wiki: Divide-and-conquer algorithm ↩︎

  3. OI Wiki: 快速傅里叶变换 ↩︎

  4. Algorithm: Fast Fourier Transform ↩︎