FFT 与 NTT

发表于 分类于 ，

概述

FFT 与 NTT 两个算法基本干的是同一件事: 将多项式由系数表示法转化为点值表示法，以及逆过程.

两种表示法之间的转化在有关生成函数的题目中随处可见， 同时也可以用来快速求离散卷积.

点值表示法

设 \(n-1\) 阶多项式 \(A(x) = \sum^{n-1}_{i=0} a_ix^i\). 这样由各项的系数组成的数列 \(a_k\) 表示了多项式 \(A(x)\).

由拉格朗日插值法或矩阵相关定理可知， n 个横坐标不同的点可以唯一确定一个阶数小于 n 的多项式，则可以用:

\[ (x_0,A(x_0)),(x_1,A(x_1))\cdots (x_n,A(x_n)) \]

这样的点值对序列来唯一确定 \(A(x)\).

这样表示的好处就是便于处理多项式之间的运算， 只需要对值进行运算即可.

FFT

使得快速傅里叶变换 (Fast Fourier Transform, FFT)“快速” 的原因在于 FFT 使用特定的一组横坐标来求对应的 \(A(x)\), 并且发掘出了其中的内在联系.

具体来讲，FFT 使用的横坐标由傅里叶变换而来， 是在复平面单位圆上的复数.

设 \(\omega_n = e^{-2\pi {\rm i}/n} = cos\frac{2\pi}{n} - {\rm i}sin\frac{2\pi}{n}\). 取 \(\omega_n^0, \omega_n^1,\cdots,\omega_n^{n-1}\) 为横坐标。由复数有关性质， 有:

\[ \omega_{2n}^{2k} = \omega_n^k \\ \omega_{2n}^{k+n} = -\omega_{2n}^{k} \]

(这两条性质画个单位圆把点标上去就懂了)

将 \(A(x)\) 按奇偶拆成两个多项式:

\[ \begin{aligned} A(x) = & a_0 + a_1x + \cdots + a_{2n-1}x^{2n-1} \\ = & (a_0 + a_2x^2 + \cdots + a_{2n-2}x^{2n-2}) \\ & + (a_1x + \cdots + a_{2n-1}x^{2n-1}) \\ = & A_0(x^2) + xA_1(x^2) \end{aligned} \]

代入 \(\omega_n^k\), 有:

\[ \begin{aligned} A(\omega_n^k) & = A_0(\omega_n^{2k}) + \omega_n^kA_1(\omega_n^{2k}) \\ & = \color{red}{A_0(\omega^k_{n/2}) + \omega^k_nA_1(\omega^k_{n/2})} \end{aligned} \]

代入 \(\omega_n^{k+(n/2)}\), 有:

\[ \begin{aligned} A(\omega_n^k) & = A_0(\omega_n^{2k+n}) + \omega_n^{k+(n/2)}A_1(\omega_n^{2k+n}) \\ & = \color{red}{A_0(\omega^k_{n/2}) - \omega^k_nA_1(\omega^k_{n/2})} \end{aligned} \]

所以只要得到所有的 \(A_0(\omega_{n/2}^k)\) 和 \(A_1(\omega_{n/2}^k)\) 便可以得到所有的 \(A(\omega_n^k)\), 而 \(A_0(\omega_{n/2}^k)\) 和 \(A_1(\omega_{n/2}^k)\) 则可以通过递归求出.

每次递归都将多项式划分成两部分，合并的时间复杂度为 \(O(n)\), 所以整个算法的时间复杂度是 \(O(n\log n)\).

注意实现时先将系数扩充为 \(2^n\) 阶， 这样更好划分多项式.

逆变换

实际上，若我们设矩阵:

\[ X = \begin{bmatrix} 1 & 1 & 1 & \cdots & 1 \\ 1 & \omega_n^1 & \omega_n^2 & \cdots & \omega_n^{n-1} \\ 1 & \omega_n^2 & \omega_n^4 & \cdots & \omega_n^{2(n-1)} \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 1 & \omega_n^{n-1} & \omega_n^{2(n-1)} & \cdots & \omega_n^{(n-1)(n-1)} \end{bmatrix} \]

则整个 FFT 可以写成:

\[ X\vec{a} = \vec{y} \]

其中 \(\vec{a}\) 表示系数的列向量， \(\vec{y}\) 表示值的列向量.

则逆变换意味着我们需要寻找 \(X\) 的逆矩阵 \(X^{-1}\).

令矩阵 B 满足 \(b_{ij} = \frac{1}{n}\omega_n^{-ij}\)(下标从 0 开始), 则矩阵 \(XB\) 的每一项满足:

\[ \begin{aligned} z_{ij} & = \frac{1}{n}\sum_{k=0}^{n-1} \omega_n^{ik}\omega_n^{-kj} \\ & = \frac{1}{n} \sum_{k=0}^{n-1} \omega_n^{k(i-j)} \\ & = \begin{cases} 1, i=j \\ \frac{1-\omega_n^{n(i-j)}}{n(1-\omega_n^{i-j})}, i \neq j \end{cases} \end{aligned} \]

当 \(i \neq j\) 时上式始终为 0, 这说明 \(B\) 确实是 \(X\) 的逆矩阵.

可以看出 \(B\) 矩阵形式上与 \(X\) 很相似， 事实上逆 FFT 只需要将求值的横坐标由 \(\omega_n^k\) 换为 \(\omega_n^{-k}\), 然后各项乘 \(1/n\) 即可.

实现代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
typedef complex<double> cpd;
cpd t[];    //临时数组
void fft(cpd * const f,int len,int iv) {    //多项式系数, 多项式阶数, 是否求逆
    if(len==1) {
        return; //递归边界
    }
    int mlen = len>>1;
    for(int i(0);i<len;i+=2) {  //划分多项式
        t[i/2] = f[i];
        t[i/2+mlen] = f[i+1];
    }
    memcpy(f,t,sizeof(cpd)*len);
    fft(f,mlen,iv); fft(f+mlen,mlen,iv);    //递归
    for(int i(0);i<mlen;++i) {
        cpd w = {cos(2.*Pi*i/len),-iv*sin(2*Pi*i/len)};  //横坐标
        t[i] = f[i] + w*f[i+mlen];
        t[i+mlen] = f[i] - w*f[i+mlen];
    }
    memcpy(f,t,sizeof(cpd)*len);
}

优化

每次向下递归都需要将多项式划分，时间开销比较大， 我们可以直接一步调整到位，然后向上回溯直接使用迭代实现. 这样可以减少很多的时间开销.

通过观察可知，每个系数 \(a_k\) 的最后位置为 \(k\) 的二进制反转.

(事实上也可以通过理解得到这个结论: 每次对奇偶项系数进行划分相当于将二进制的最后一位提到最高位后重新排序)

同时为了提高精度并且减小时间开销， 可以进行单位根预处理 , 预处理出最小单位根的各个次方. (事实上精度是 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
39
40
41
42
43
typedef complex<double> cpd;
int len;    //扩充后长度
int r[];    //位倒置后的数
cpd pw[Mn]; //单位根预处理
void init() {   //初始化
    int hb(-1);
    for(len=1; len<=n+m; len<<=1, ++hb) ;   //扩充长度
    for(int i(1);i<len;++i) {
        r[i] = (r[i>>1]>>1) | ((i&1)<<hb);  //位倒置
    }
    /*单位根预处理*/
    pw[0] = cpd(1,0);
    for(int i(1);i<len;++i) {
        pw[i] = polar(1.L,-Pi*i/len);
    }
}

void fft(cpd * const f,int inv) {   //多项式, 是否求逆
    if(inv==-1) {
        //若求逆则直接将数组倒置, 这样就不用多算一次单位根了
        for(int i(1);i<(len>>1);++i) {
            swap(f[i],f[len-i]);
        }
    }
    for(int i(1);i<len;++i) {   //一步调整
        if(i<r[i]) {        //不加这个if则会交换两次
            swap(f[i],f[r[i]]);
        }
    }
    for(int ml(1);ml<len;ml<<=1) {  //当前处理长度的一半
        for(int p(0);p<len;p+=(ml<<1)) {    //当前处理位置起点
            for(int i(p);i<(p|ml);++i) {
                cpd w = pw[(i^p)*(len/ml)]; //取横坐标
                /*
                一个小优化: 减少复数乘法的运算次数
                */ 
                cpd x = f[i], y = w*f[i|ml];
                f[i] = x+y;
                f[i|ml] = x-y;
            }
        }
    }
}

NTT

数论变换 (Number Theory Transform, NTT) 的思想与 FFT 差不多， 区别在于加上取模的操作后可以利用数论的一些概念来代替复数.

首先需要介绍” 原根” 的概念:

设 \(m\) 是正整数，\(a\) 是整数，若 \(a\) 模 \(m\) 的阶等于 \(\phi(m)\)，则称 \(a\) 为模 \(m\) 的一个原根.

说人话就是对于所有 \([1,\phi(m)]\) 的两个不同整数 \(i,j\), 有 \(a^i \not\equiv a^j(mod\ m)\).

设 \(g\) 为素数 \(P\) 的一个原根，令 \(\omega_n=g^{(P-1)/n}\ mod\ P\), 则由原根的性质，可以得到一些与 FFT 相似的性质:

\[ \omega_{2n}^{2k} \equiv \omega_{n}^k (mod\ P)\tag{1} \]

这点十分显然，分数上下约分即可.

\[ \omega_{2n}^{k+n} \equiv -\omega_{2n}^{k}(mod\ P)\tag{2} \]

等同于证明 \(g^{(P-1)/2}\equiv -1(mod\ P)\). 由数论知识可知 \(g^{(P-1)/2}\) 只有可能是 \(\pm 1\), 而若其等于 1 则与原根的定义相矛盾 (\(g^{P-1}\) 由费马小定理可知等于 1), 所以 \(g^{(P-1)/2}\equiv -1(mod\ P)\).

根据这两条性质在 \(mod\ P\) 的条件下可以令 \(\omega_n=g^{(P-1)/n}\ mod\ P\) 进行变换.

需要注意的是 \(\frac{P-1}{n}\) 必须是整数，所以 \(P\) 需要满足 \(k \times 2^m+1\) 的形式，同时中间的 \(m\) 需要足够大使得 \(2^m \geq n\).

常见的模数有:

\[ \begin{aligned} 998244353 & = 119 \times 2^{23} + 1 \\ 1004535809 & = 479 \times 2^{21} + 1 \\ 469762049 & = 7 \times 2^{26} + 1 \end{aligned} \]

它们都有一个原根 \(3\).

实现代码:

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
typedef long long LL;
const LL MP(7*17*(1<<23)+1);    //998244353
const LL G(3), Gi(332748118);    //原根及其逆元
LL ksm(LL a,LL b) {    //快速幂
    LL ret(1),tmp(a%MP);
    while(b) {
        if(b&1) {
            ret = ret*tmp%MP;
        }
        tmp = tmp*tmp%MP;
        b >>= 1;
    }
    return ret;
}
void ntt(LL * const f,int len,bool iv) {    //同FFT
    for(int i(1);i<len;++i) {
        if(i<r[i]) {
            swap(f[i],f[r[i]]);
        }
    }
    for(int l(2);l<=len;l<<=1) {
        LL wn = ksm(iv?Gi:G,(MP-1)/l);
        int ml = l>>1;
        for(int p(0);p<len;p+=l) {
            LL w = 1;
            for(int i(p);i<p+ml;++i, w=w*wn%MP) {	//注意取模
                LL x = f[i], y = w*f[i|ml]%MP;
                f[i] = (x+y)%MP;
                f[i|ml] = (x-y+MP)%MP;
            }
        }
    }
}

用途

卷积

设两个多项式:

\[ A(x) = \sum_{i=0}^{n-1} a_ix^i \\ B(x) = \sum_{i=0}^{m-1} b_ix^i \]

将两个多项式相乘可得:

\[ \begin{aligned} & A(x) \times B(x) \\ = & \sum_{i=0}^{n+m-2} (\sum_{j=0}^i a_jb_{i-j}) x^i \end{aligned} \]

注意到括号里的系数便是两个多项式系数组成的数列的卷积. 可以使用 FFT 将两个多项式变成相同的点值表示， 值相乘后再逆变换便得到了所有位置的卷积。时间复杂度为 \(O(n\log n)\)

大数乘法

事实上大数乘法是卷积的特殊形式，只需要卷积后再处理进位即可.