FFT & NTT - 快速傅里叶变换 & 快速数论变换

更好的阅读体验戳此进入

（建议您从上方链接进入我的个人网站查看此 Blog，在 Luogu 中图片会被墙掉，部分 Markdown 也会失效）

写在前面

该博客仅为记录学习中的笔记及个人理解，不保证正确性，同时欢迎各位纠正。

图片没有放在图床上，全都是丢在自己的网站上，带宽较低可能加载较慢。

目的

$C(x) = A(x) \ast B(x)$ ，或者表达为

c (i) = \sum_{j = 0}^{i} a (j) \times b (i - j)

$a(i), b(i), c(i)$ $A(x), B(x), C(x)$ 为多项式 )

前置知识

原根

详细定义可参考知乎 OI-WIKI $m$ $a$ $a^1, a^2, \cdots, a^{\varphi(m)} \bmod m$ $\varphi(m)$ 表示欧拉函数。

单位根

$\epsilon^n = 1 (\epsilon \neq 1)$ $\epsilon$ $n$ 次单位根，其可以为模意义下的或复数意义下的。

模意义下的（原根）

$m$ $m$ $g$ $\gcd(g, m) = 1$ $n \mid m - 1$ $\epsilon = g^{\frac{m - 1}{n}}$ 。

$d \le m - 2$ $\epsilon^d \bmod{m}$ 各不相同。

证明：

\begin{aligned} (g^{\frac{m - 1}{n}})^{n} & \equiv 1 (\mod m) \\ ⟺ & g^{m - 1} & \equiv 1 (\mod m) \end{aligned}

由费马小定理可知显然成立

复数意义下的

$(1, 0)$ $\epsilon_n^k$ $\epsilon_n^1$ $(\epsilon_n^1)^k = \epsilon_n^k$ $\epsilon_n^1$ $n$ 次单位根。

$\omega$ $\epsilon$ ，仅为表示方式的区别而已。

单位根性质

$n$ 次单位根有如下式子

(ϵ^{k})^{2} = (ϵ^{k + \frac{n}{2}})^{2}

证明

\begin{aligned} (ϵ^{k + \frac{n}{2}})^{2} & = (ϵ^{k})^{2} \times ϵ^{2} \\ = (ϵ^{k})^{2} \end{aligned}

单位根求法

复数意义下

$\omega_n^k$ $\omega_n^k = cos(2\pi \div n \times k) + sin(2\pi \div n \times k)i$ 。

模意义下的（原根）

$d \le m - 2$ $d \equiv 1 (\bmod{m})$ 则其为原根。

$m - 1$ $p_i$ $\epsilon$ $\epsilon^{\frac{m - 1}{p_i}} \equiv 1 (\bmod{m})$ $\epsilon$ 为原根。

等比数列求和公式

详细证明

S_{n} = a_{1} \frac{1 - q^{n}}{1 - q}

正文

单位根反演

$n$ $\epsilon$ $a_1 = \epsilon^0 = 1$ $q = \epsilon^v$ ，的等比数列的求和为

\begin{aligned} \sum_{i = 0}^{n - 1} ϵ^{v i} & = \frac{1 - ϵ^{n v}}{1 - ϵ^{v}} (ϵ^{v} \neq 1) \\ = 0 \end{aligned}

$n$ $\epsilon^v = 1$ $= 1$ 。

且又有如下式子

ϵ^{v} = 1 ⟺ n ∣ v

综上则有如下式子

\begin{matrix} \frac{1}{n} \sum_{i = 0}^{n - 1} ϵ^{v i} = {\begin{cases} 0 & n ∤ v \\ 1 & n ∣ v \end{cases} \end{matrix}

此即为单位根反演

推式子

$v = p + q - i$ $n \mid v$ ，

且令

\begin{matrix} d (x) = {\begin{cases} 0 & p + q ≢ i mod n \\ 1 & p + q \equiv i mod n \end{cases} \end{matrix}

显然有如下式子

\begin{aligned} c (i) & = \sum_{j = 0}^{i} a (j) \times b (i - j) \\ = \sum_{p} \sum_{q} a (p) \times b (q) \times d (x) \\ = \sum_{p = 0}^{n - 1} \sum_{q = 0}^{n - 1} a (p) \times b (q) \times \frac{1}{n} \sum_{k = 0}^{n - 1} ϵ^{k \times (p + q - i)} \\ = \sum_{p = 0}^{n - 1} a (p) \sum_{q = 0}^{n - 1} b (q) \times \frac{1}{n} \sum_{k = 0}^{n - 1} ϵ^{k p} ϵ^{k q} ϵ^{- k i} \\ = \sum_{p = 0}^{n - 1} a (p) ϵ^{k p} \sum_{q = 0}^{n - 1} b (q) ϵ^{k q} \times \frac{1}{n} \sum_{k = 0}^{n - 1} ϵ^{- k i} \end{aligned}

⟺ n \times c (i) \sum_{k = 0}^{n - 1} ϵ^{k i} = \sum_{p = 0}^{n - 1} a (p) ϵ^{k p} \sum_{q = 0}^{n - 1} b (q) ϵ^{k q}

观察最后两个式子，可以发现如下两个式子

c (i) = \sum_{p = 0}^{n - 1} a (p) ϵ^{k p} \sum_{q = 0}^{n - 1} b (q) ϵ^{k q} \times \frac{1}{n} \sum_{k = 0}^{n - 1} ϵ^{- k i}

⟺ n \times c (i) \sum_{k = 0}^{n - 1} ϵ^{k i} = \sum_{p = 0}^{n - 1} a (p) ϵ^{k p} \sum_{q = 0}^{n - 1} b (q) ϵ^{k q}

$(\epsilon^p, f(p))$ $p$ $g(p)$ ，有

\begin{aligned} f (p) & = \sum_{i = 0}^{n - 1} ϵ^{p i} g (i) \\ g (p) & = \frac{1}{n} \sum_{i = 0}^{n - 1} ϵ^{- p i} f (i) \end{aligned}

证明

/ / T O D O

$f(p)$ $g(p)$ 的过程即为IDFT。

由定义显然有

D F T (C, i) = D F T (A, i) \times D F T (B, i)

$A, B, C$ 均代表该多项式）

又有

I D F T (D F T (C)) = C

证明

/ / T O D O

$O(n^2)$

继续推式子

对于

f (p) = \sum_{i = 0}^{n - 1} ϵ^{p i} g (i)

$f(p)$ $f(p + 2^k)$ $n = 2^{k + 1}$

且令

\begin{matrix} d_{1} (x) = {\begin{cases} 0 & i \equiv 0 (mod 2) \\ 1 & i \equiv 1 (mod 2) \end{cases} \\ d_{2} (x) = {\begin{cases} 1 & i \equiv 0 (mod 2) \\ 0 & i \equiv 1 (mod 2) \end{cases} \end{matrix}

由单位根的性质可以得到以下式子

\begin{aligned} f (p + 2^{k}) & = \sum_{i = 0}^{n - 1} (ϵ^{p + 2^{k}})^{i} g (i) \\ = \sum_{i = 0}^{n - 1} (ϵ^{p})^{i} g (i) d_{1} (i) + \sum_{i = 0}^{n - 1} (ϵ^{p})^{i} u g (i) d_{2} (i) \end{aligned}

$i$ $0$ 。

\sum_{i = 0}^{n - 1} (ϵ^{p})^{i} u g (i) d (i)

$i$ 为奇数时，可以有如下推导

\begin{aligned} \sum_{i = 0}^{n - 1} (ϵ^{p + 2^{k}})^{i} g (i) & = \sum_{i = 0}^{n - 1} (ϵ^{p})^{i} ϵ^{2^{k}} g (i) \end{aligned}

此时显然有

\begin{aligned} (ϵ^{2^{k}})^{2} & = ϵ^{2^{k + 1}} \end{aligned}

$\epsilon$ $2^{k + 1}$ 次单位根，所以显然有

ϵ^{2^{k + 1}} = 1

$\epsilon^{2^k}$ $u$ 代替。

此时可以考虑令

\begin{aligned} f^{^{'}} (p) & = & \sum_{i = 0}^{n - 1} (ϵ^{p})^{i} g (i) d_{1} (i) \\ f^{^{'}} (p + 2^{k}) & = & \sum_{i = 0}^{n - 1} (ϵ^{p})^{i} g (i) d_{2} (i) \end{aligned}

$\epsilon$ $\dfrac{n}{2}$ 次单位根）

\begin{aligned} f^{^{'}} (p) & = \sum_{i = 0}^{\frac{n}{2} - 1} (ϵ^{p})^{i} g (i \times 2) \\ f^{^{'}} (p + 2^{k}) & = \sum_{i = 0}^{\frac{n}{2} - 1} (ϵ^{p})^{i} g (i \times 2 + 1) \end{aligned}

再将式子转化为

\begin{aligned} f^{^{'}} (p) & = \sum_{i = 0}^{n^{^{'}}} (ϵ^{p})^{i} g_{1} (i) \\ f^{^{'}} (p + 2^{k}) & = \sum_{i = 0}^{n^{^{'}}} (ϵ^{p})^{i} g_{2} (i) \end{aligned}

$n = 2$ 时进行回溯。

Code


xxxxxxxxxx
98
1
#define _USE_MATH_DEFINES
2
#include <bits/stdc++.h>
3

4
#define PI M_PI
5
#define E M_E
6
#define DFT true
7
#define IDFT false
8
#define eps 1e-6
9

10
#define comp complex < double >
11

12
/******************************
13
abbr
14
pat -> pattern
15
pol/poly -> polynomial
16
omg -> omega
17
******************************/
18

19
using namespace std;
20

21
mt19937 rnd(random_device{}());
22
int rndd(int l, int r){return rnd() % (r - l + 1) + l;}
23

24
typedef unsigned int uint;
25
typedef unsigned long long unll;
26
typedef long long ll;
27

28
class Polynomial{
29
    private:
30
        int lena, lenb;
31
        int len;
32
        comp A[1100000], B[1100000];
33
    public:
34
        comp Omega(int, int, bool);
35
        void Init(void);
36
        void FFT(comp*, int, bool);
37
        void MakeFFT(void);
38
}poly;
39

40
template<typename T = int>
41
inline T read(void);
42

43
int main(){
44
    poly.Init();
45
    poly.MakeFFT();
46

47
    fprintf(stderr, "Time: %.6lf\n", (double)clock() / CLOCKS_PER_SEC);
48
    return 0;
49
}
50
void Polynomial::MakeFFT(void){
51
    FFT(A, len, DFT), FFT(B, len, DFT);
52
    for(int i = 0; i <= len; ++i)A[i] *= B[i];
53
    FFT(A, len, IDFT);
54
    for(int i = 0; i <= lena + lenb - 2; ++i)
55
        printf("%d%c", int(A[i].real() / len + eps + 0.5), i == lena + lenb - 1 ? '\n' : ' ');
56
}
57
void Polynomial::FFT(comp* pol, int len, bool pat){
58
    if(len == 1)return;
59
    comp sA[len / 2 + 10], sB[len / 2 + 10];
60
    for(int i = 0; i <= len / 2 - 1; ++i){
61
        sA[i] = pol[i * 2];
62
        sB[i] = pol[i * 2 + 1];
63
    }
64
    FFT(sA, len / 2, pat), FFT(sB, len / 2, pat);
65
    for(int i = 0; i <= len / 2 - 1; ++i){
66
        comp omg = Omega(len, i, pat);
67
        pol[i] = sA[i] + omg * sB[i];
68
        pol[i + len / 2] = sA[i] - omg * sB[i];
69
    }
70
}
71
void Polynomial::Init(void){
72
    lena = read(), lenb = read();
73
    for(int i = 0; i <= lena; ++i)A[i].real((double)read());
74
    for(int i = 0; i <= lenb; ++i)B[i].real((double)read());
75
    len = 1;
76
    lena++, lenb++;
77
    while(len <= lena + lenb)len <<= 1;
78
}
79
comp Polynomial::Omega(int n, int k, bool pat){
80
    if(pat == DFT)return comp(cos(2 * PI * k / n), sin(2 * PI * k / n));
81
    return conj(comp(cos(2 * PI * k / n), sin(2 * PI * k / n)));
82
}
83

84
template<typename T>
85
inline T read(void){
86
    T ret(0);
87
    short flag(1);
88
    char c = getchar();
89
    while(c != '-' && !isdigit(c))c = getchar();
90
    if(c == '-')flag = -1, c = getchar();
91
    while(isdigit(c)){
92
        ret *= 10;
93
        ret += int(c - '0');
94
        c = getchar();
95
    }
96
    ret *= flag;
97
    return ret;
98
}

优化

洛谷模板 $\text{Cooley - Tukey}$ 算法。

首先观察如下递归过程（图片来源）

FFT_1

通过观察我们即可发现（这真是人类能想出来的吗）对于每一个数的位置，显然是进行了一次二进制的反转，如 1 的位置从 001 变成了 100，那么我们便可以利用这个性质对位置进行反转。

这里提供两种写法

$O(nlogn)$


xxxxxxxxxx
9
1
int size(0);
2
while((1 << size) < len - 1)++size;
3
for(int i = 0; i <= len - 1; ++i){
4
    int tmp(0);
5
    for(int j = 0; j <= size; ++j){
6
        if((1 << j) & i) tmp |= (1 << (size - j - 1));
7
    }
8
    if(i < tmp)swap(pol[i], pol[tmp]);
9
}

$i < tmp$ 的判断是为了使其只会交换一次。

$O(n)$


xxxxxxxxxx
7
1
int pos[len + 10];
2
memset(pos, 0, sizeof(pos));
3
for(int i = 0; i < len; ++i){
4
    pos[i] = pos[i >> 1] >> 1;
5
    if(i & 1)pos[i] |= len >> 1;
6
}
7
for(int i = 0; i < len; ++i)if(i < pos[i])swap(pol[i], pol[pos[i]]);

这种方法我就不严格地证明了（主要我也不会），就从找规律的角度来研究一下这个线性递推的式子。

$0101110$ $0010111$ $1110100$ $0111010$ $1$ $0$ 即可。

对于 Reverse 后合并的过程显然我们可以通过从倒数第二层开始，模拟递归形式的操作，这部分较为显然便不再赘述。

值得注意的一个点是当我们更新数组时，由于非递归写法，可能会对需要用到的变量进行覆盖，所以这时我们显然可以将原数组复制一份，这样的空间时可以接受的，当然更好的做法就是将会被覆盖的那个变量存起来再进行操作，如下。


xxxxxxxxxx
11
1
Reverse(pol, len);
2
for(int size = 2; size <= len; size <<= 1){
3
    for(comp* p = pol; p != pol + len; p += size){
4
        int mid(size >> 1);
5
        for(int i = 0; i < mid; ++i){
6
            auto tmp = Omega(size, i, pat) * p[i + mid];
7
            p[i + mid] = p[i] - tmp;
8
            p[i] = p[i] + tmp;
9
        }
10
    }
11
}

最后贴上优化后的完整代码


xxxxxxxxxx
108
1
#define _USE_MATH_DEFINES
2
#include <bits/stdc++.h>
3
#include <mmintrin.h>
4

5
#define PI M_PI
6
#define E M_E
7
#define DFT true
8
#define IDFT false
9
#define eps 1e-6
10

11
#define comp complex < double >
12

13
/******************************
14
abbr
15
pat -> pattern
16
pol/poly -> polynomial
17
omg -> omega
18
******************************/
19

20
using namespace std;
21

22
mt19937 rnd(random_device{}());
23
int rndd(int l, int r){return rnd() % (r - l + 1) + l;}
24

25
typedef unsigned int uint;
26
typedef unsigned long long unll;
27
typedef long long ll;
28

29
class Polynomial{
30
    private:
31
        int lena, lenb;
32
        int len;
33
        comp A[2100000], B[2100000];
34
    public:
35
        comp Omega(int, int, bool);
36
        void Init(void);
37
        void FFT(comp*, int, bool);
38
        void Reverse(comp*);
39
        void MakeFFT(void);
40
}poly;
41

42
template<typename T = int>
43
inline T read(void);
44

45
int main(){
46
    poly.Init();
47
    poly.MakeFFT();
48

49
    fprintf(stderr, "Time: %.6lf\n", (double)clock() / CLOCKS_PER_SEC);
50
    return 0;
51
}
52
void Polynomial::MakeFFT(void){
53
    FFT(A, len, DFT), FFT(B, len, DFT);
54
    for(int i = 0; i <= len; ++i)A[i] *= B[i];
55
    FFT(A, len, IDFT);
56
    for(int i = 0; i <= lena + lenb - 2; ++i)
57
        printf("%d%c", int(A[i].real() / len + eps + 0.5), i == lena + lenb - 2 ? '\n' : ' ');
58
}
59
void Polynomial::Reverse(comp* pol){
60
    int pos[len + 10];
61
    memset(pos, 0, sizeof(pos));
62
    for(int i = 0; i < len; ++i){
63
        pos[i] = pos[i >> 1] >> 1;
64
        if(i & 1)pos[i] |= len >> 1;
65
    }
66
    for(int i = 0; i < len; ++i)if(i < pos[i])swap(pol[i], pol[pos[i]]);
67
}
68
void Polynomial::FFT(comp* pol, int len, bool pat){
69
    Reverse(pol);
70
    for(int size = 2; size <= len; size <<= 1){
71
        for(comp* p = pol; p != pol + len; p += size){
72
            int mid(size >> 1);
73
            for(int i = 0; i < mid; ++i){
74
                auto tmp = Omega(size, i, pat) * p[i + mid];
75
                p[i + mid] = p[i] - tmp;
76
                p[i] = p[i] + tmp;
77
            }
78
        }
79
    }
80
}
81
void Polynomial::Init(void){
82
    lena = read(), lenb = read();
83
    for(int i = 0; i <= lena; ++i)A[i].real((double)read());
84
    for(int i = 0; i <= lenb; ++i)B[i].real((double)read());
85
    len = 1;
86
    lena++, lenb++;
87
    while(len <= lena + lenb)len <<= 1;
88
}
89
comp Polynomial::Omega(int n, int k, bool pat){
90
    if(pat == DFT)return comp(cos(2 * PI * k / n), sin(2 * PI * k / n));
91
    return conj(comp(cos(2 * PI * k / n), sin(2 * PI * k / n)));
92
}
93

94
template<typename T>
95
inline T read(void){
96
    T ret(0);
97
    short flag(1);
98
    char c = getchar();
99
    while(c != '-' && !isdigit(c))c = getchar();
100
    if(c == '-')flag = -1, c = getchar();
101
    while(isdigit(c)){
102
        ret *= 10;
103
        ret += int(c - '0');
104
        c = getchar();
105
    }
106
    ret *= flag;
107
    return ret;
108
}

NTT

前面我们已知 FFT 是在复数意义下利用单位复根的性质进行优化，而 NTT 则是在模意义下的，对于模意义下的单位根替代品则为原根，至于证明这里不再赘述，可以在此处查看。

$9 \times 9 \times 10^6$ $998244353$ $3$ ，如何求原根可以在前置知识中找到。

$\omega_n$ $g^{\frac{MOD - 1}{size}}$ $A(i) \div len$ $A(i) \times inv(len) \bmod{MOD}$ 。

Code:


xxxxxxxxxx
116
1
#define _USE_MATH_DEFINES
2
#include <bits/stdc++.h>
3
#include <mmintrin.h>
4

5
#define PI M_PI
6
#define E M_E
7
#define DFT true
8
#define IDFT false
9
#define eps 1e-6
10
#define MOD 998244353
11

12
/******************************
13
abbr
14
pat -> pattern
15
pol/poly -> polynomial
16
******************************/
17

18
using namespace std;
19

20
mt19937 rnd(random_device{}());
21
int rndd(int l, int r){return rnd() % (r - l + 1) + l;}
22

23
typedef unsigned int uint;
24
typedef unsigned long long unll;
25
typedef long long ll;
26

27
ll kpow(int a, int b){
28
    ll ret(1ll), mul((ll)a);
29
    while(b){
30
        if(b & 1)ret = (ret * mul) % MOD;
31
        b >>= 1;
32
        mul = (mul * mul) % MOD;
33
    }
34
    return ret;
35
}
36
class Polynomial{
37
    private:
38
        int lena, lenb;
39
        int len;
40
        int g, inv_g;
41
        int A[2100000], B[2100000];
42
    public:
43
        int Omega(int, int, bool);
44
        void Init(void);
45
        void NTT(int*, int, bool);
46
        void Reverse(int*);
47
        void MakeNTT(void);
48
}poly;
49

50
template<typename T = int>
51
inline T read(void);
52

53
int main(){
54
    poly.Init();
55
    poly.MakeNTT();
56
    fprintf(stderr, "Time: %.6lf\n", (double)clock() / CLOCKS_PER_SEC);
57
    return 0;
58
}
59
void Polynomial::MakeNTT(void){
60
    NTT(A, len, DFT), NTT(B, len, DFT);
61
    for(int i = 0; i <= len; ++i)A[i] = ((ll)A[i] * B[i]) % MOD;
62
    NTT(A, len, IDFT);
63
    int mul_inv = kpow(len, MOD - 2);
64
    for(int i = 0; i <= lena + lenb - 2; ++i)
65
        printf("%d%c", (ll)A[i] * mul_inv % MOD, i == lena + lenb - 2 ? '\n' : ' ');
66
}
67
void Polynomial::Reverse(int* pol){
68
    int pos[len + 10];
69
    memset(pos, 0, sizeof(pos));
70
    for(int i = 0; i < len; ++i){
71
        pos[i] = pos[i >> 1] >> 1;
72
        if(i & 1)pos[i] |= len >> 1;
73
    }
74
    for(int i = 0; i < len; ++i)if(i < pos[i])swap(pol[i], pol[pos[i]]);
75
}
76
void Polynomial::NTT(int* pol, int len, bool pat){
77
    Reverse(pol);
78
    for(int size = 2; size <= len; size <<= 1){
79
        int gn = kpow(pat == DFT ? g : inv_g, (MOD - 1) / size);
80
        for(int* p = pol; p != pol + len; p += size){
81
            int mid(size >> 1);
82
            int g(1);
83
            for(int i = 0; i < mid; ++i, g = ((ll)g * gn) % MOD){
84
                auto tmp = ((ll)g * p[i + mid]) % MOD;
85
                p[i + mid] = (p[i] - tmp + MOD) % MOD;
86
                p[i] = (p[i] + tmp) % MOD;
87
            }
88
        }
89
    }
90
}
91
void Polynomial::Init(void){
92
    lena = read(), lenb = read();
93
    for(int i = 0; i <= lena; ++i)A[i] = read();
94
    for(int i = 0; i <= lenb; ++i)B[i] = read();
95
    len = 1;
96
    lena++, lenb++;
97
    while(len < lena + lenb)len <<= 1;
98
    g = 3;
99
    inv_g = kpow(g, MOD - 2);
100
}
101

102
template<typename T>
103
inline T read(void){
104
    T ret(0);
105
    short flag(1);
106
    char c = getchar();
107
    while(c != '-' && !isdigit(c))c = getchar();
108
    if(c == '-')flag = -1, c = getchar();
109
    while(isdigit(c)){
110
        ret *= 10;
111
        ret += int(c - '0');
112
        c = getchar();
113
    }
114
    ret *= flag;
115
    return ret;
116
}

合并DFT优化

这个单独再写一个 Blog 吧，戳此进入。

写在后面

写完之后发现似乎依然没有很清晰的弄明白，然后发现有几个Blog写的更清晰易懂

一小时学会快速傅里叶变换（Fast Fourier Transform）

小学生都能看懂的FFT！！！

至于几个TODO等以后再慢慢填坑吧

UPD

update-2022_08_10 初稿

update-2022_08_17 改了一下 latex 在 cnblog 里渲染异常的问题（ luogu 里还是炸了，以后再改）

update-2022_08_17 修复 latex 在 luogu 里渲染异常的问题

update-2022_08_22 修复 latex 在 cnblog 里仍然存在的渲染异常问题

update-2022_08_22 添加了递归版程序中的 code

update-2022_08_22 进行一些小优化

update-2022_08_22 添加了非循环写法的讲解与 code

update-2022_08_22 添加了 NTT 的讲解与 code

update-2022_08_22 完善了对模意义下单位根的求法

update-2022_08_23 更改标题

update-2022_08_23 添加几个链接

update-2022_08_25 更新标题和链接