再探快速傅里叶变换(FFT)学习笔记(其二)(NTT)

再探快速傅里叶变换(FFT)学习笔记(其二)(NTT)

写在前面

为了不使篇幅过长，预计将把学习笔记分为四部分:

DFT,IDFT,FFT的定义,实现与证明:快速傅里叶变换(FFT)学习笔记(其一)
NTT的实现与证明,任意模数NTT
FFT的优化技巧
多项式相关操作

一些约定

\([p(x)]=\begin{cases}1,p(x)为真 \\ 0,p(x)为假 \end{cases}\)
本文中序列的下标从0开始
若\(s\)是一个序列,\(|s|\)表示\(s\)的长度
若大写字母如\(F(x)\)表示一个多项式,那么对应的小写字母如\(f\)表示多项式的每一项系数，即\(F(x)=\sum_{i=0}^{n-1} f_ix^i\)
由于本文主要讨论数论知识，本文中的未知数如未经说明，均为正整数。

前置知识

同余类和剩余系

定义1.1:若整数\(a\)和整数\(b\)除以正整数\(m\)的余数相等，则称\(a,b\)模\(m\)同余，记作\(a \equiv b (\bmod m)\)

定义1.2:把全体整数按照模\(m\)的余数分为若干个互不相交的集合,每一个集合称为模\(m\)的同余类(剩余类)(residue(congruence) class).余数为\(a\)的同余类记作\(\bar{a}\)

对于\(\forall a \in [0,m]\),\(\bar{a}=\{a+km \}(k \in \mathbb{Z})\)

同余类中的每个元素都可以拿来代表该同余类，称为该同余类的代表数

定义1.3:模\(m\)的同余类一共有\(m\)个,分别为\(\bar{0},\bar{1},\dots \bar{m-1}\).在它的每一个同余类中取一个元素作为代表，所有这些代表元素组成的集合称为\(m\)的完全剩余系(complete residue system)

定义1.4 \([1,m]\)中与\(m\)互质的数代表的同余类组成\(m\)的简化(既约)剩余系(reduced residue system)

例如,10的简化剩余系为\(\{\bar{1},\bar{3},\bar{7},\bar{9}\}\)

定理1.1 \(m\)的简化剩余系关于模\(m\)乘法封闭

证明:若\(\forall a,b \in [1,m]\)与\(m\)互质,那么\(ab\)也不可能与\(m\)含有相同的质因子。根据辗转相除法我们知道\(\gcd(ab \bmod m,m)=\gcd(ab,m)=1\),也就是说\(ab \bmod m\)也属于\(m\)的简化剩余系

欧拉定理

我们知道\(\varphi(m)\)表示\([1,m]\)中与\(m\)互质的b数的个数。那么易得\(m\)的简化剩余系的大小为\(\varphi(m)\)

定理2.1 若\(\gcd(a,n)=1\),则\(a^{\varphi(n)}=1 (\bmod n)\)

证明:设\(n\)的简化剩余系为\(\{\bar{a_1},\bar{a_2},\dots \bar{a_{\varphi(n)}} \}\)

\(\forall a_i,a_j,a_i \neq a_j (\bmod n)\)

又因为\(\gcd(a,n)=1\), 故\(a(a_i-a_j) \neq 0 (\bmod n)\),即\(aa_i\neq aa_j(\bmod n)\)

又因为简化剩余系关于模\(n\)乘法封闭,故\(\bar{aa_i}\)也在\(n\)的简化剩余系中。因此\(\{\bar{a_1},\bar{a_2},\dots \bar{a_{\varphi(n)}} \}\)和\(\{\bar{aa_1},\bar{aa_2},\dots \bar{aa_{\varphi(n)}} \}\)都可以表示\(n\)的简化剩余系(可以理解为对应顺序不同).所以有

\(a^{\varphi(n)}a_1a_2\dots a_{\varphi(n)} \equiv a_1a_2\dots a_{\varphi(n)} (\bmod n)\)

所以\(a^{\varphi(n)} \equiv 1 (\bmod n)\)

定理2.2,若\(p\)为质数且\(\gcd(a,p)=1\),\(a^{p-1} \equiv 1 (\bmod p)\)

由定理2.1显然

定理2.1被称为欧拉定理,而定理2.2被称为费马小定理

阶

根据欧拉定理我们知道,存在正整数\(d \leq m-1\),使得\(a^d \equiv 1 (\bmod m)\).(显然\(\varphi(m) \leq m-1\),\(m\)为素数时取等).但是,最小的\(d\)是多少呢?

定义3.1:设\(m>1\),且\(\gcd(a,m)=1\),那么使得\(a^r \equiv 1 (\bmod m)\)的最小正整数\(r\)称为\(a\)对模\(m\)的阶(指数)(order),记为\(\delta_m(a)\)

定理3.1: 对于正整数\(d>1\),\(a^d \equiv 1 (\bmod m)\)的充要条件是\(\delta_m(a)|d\).

证明:

由于\(m>1,\gcd(a,m)=1\),易得\(\forall j \geq 0,a^j \bmod m \ne 0\)(\(a\)与\(m\)没有公共质因数,那么\(a^j\)也一定和\(m\)没有公共的质因数).

不妨设\(a^j=q_jm+r_j(0<r_j<m)\)

这样,取\(j=0,1,2 \dots m-1\),其中\(m\)个余数\(r_0,r_1,\dots r_{m-1}\)最多取\(m-1\)个值,根据抽屉原理,必有两个相等,记为\(r_i,r_k\).不妨设\(0 \leq k<i<m\),那么\(m(q_i-q_k)=a^i-a^k=a^k(a^{i-k}-1)\). 所以\(m|a^k(a^{i-k}-1)\)

又因为\(\gcd(a^k,m)=1\),所以把互质的因数去掉也不影响整除性,所以\(m|(a^{i-k}-1)\)，即\(a^{i-k} \equiv 1 (\bmod m)\).取\(d=i-k\)就证明了该定理.

定理3.2:设\(k\)是非负整数,则有

\[\delta_m(a_k)=\frac{\delta_m(a)}{\gcd(\delta_m(a),k)}\]

设\(\delta=\delta_m(a),\delta_1=\frac{\delta}{\gcd(\delta,k)},\delta_2=\delta_m(a^k)\),那要证明的就是\(\delta_1=\delta_2\)

显然\(k\delta_1=\text{lcm}(\delta,k)\),因此\(\delta|k\delta_1\), 根据定理3.1\(a^{k\delta_1}=(a^k)^{\delta_1} \equiv 1 (\bmod m)\).对于\(a^k\),根据原根的最小性和定理3.1有\(\delta_2|\delta_1\)

又根据原根的定义\(a^{k\delta_2}\equiv 1(\bmod m)\),对于\(a\)来说,\(\delta| k\delta_2\).设,则\(\frac{\delta}{\gcd(\delta,k)}=\frac{k\delta_2}{\gcd(\delta,k)}\),注意到左边就是\(\delta_2\)所以\(\delta_1|\delta_2\)

原根

定义4.1:设\(m \geq 1,\gcd(a,m)=1\),若\(\delta_m(a)=\varphi(m)\),则称\(a\)为\(m\)的原根

定理4.1:若\(m\)有原根\(g\),模\(m\)的简化剩余系可以表示为\(\{\bar{g^0},\bar{g^1},\dots \bar{g^{\varphi(m)-1}} \}\)

证明:

引理4.1.1: \(g^0,g^1 \dots g^{\varphi(m)-1}\)两两模\(m\)不同余

反证法,假设存在\(0 \leq i<j \leq \varphi(m)-1\),使得\(g^i \equiv g^j (\bmod m)\),那么\(g^{j-i} \equiv 1(\bmod m)\),即\(j-i< \varphi(m)\),这与原根的最小性矛盾

引理4.1.2:\(g^0,g^1 \dots g^{\varphi(m)-1}\)均与\(m\)互质

由于\(\gcd(g,m)=1\),所以\(g^k \bmod m \neq 0\)

结合两个引理和简化剩余系的定义，我们证明了定理4.1

定理4.1的推论:若\(p\)是素数,\(g^0,g^1 \dots g^{p-2}\)恰好构成\(1,2,\dots p-1\)的一个排列

证明:\(\varphi(p)=p-1\),显然得证

定理4.2 \(m\)的原根有\(\varphi(\varphi(m))\)个

证明:根据定理4.2,\(g^0,g^1 \dots g^{\varphi(m)-1}\)遍历\(m\)的简化剩余系,也就是说我们只需要考虑\(g\)的幂\(g^k\).

若\(a^k\)是\(m\)的原根,\(\delta_m(a^k)=\varphi(m)=\delta_m(a)\).

根据定理3.2,\(\delta_m(a^k)=\frac{\delta_m(a)}{\gcd(\delta_m(a),k)}\)

那么\(\gcd(\delta_m(a),k)=1\),即\(\gcd(\varphi(m),k)=1\).满足条件的\(k\)的个数就是\(\varphi(\varphi(m))\)

求原根

定理4.3:对于任意正整数\(m \geq 2\),\(\varphi(m)\)的所有的不同的素因数为\(q_1,q_2 \dots q_s\),那么\(g\)是模\(p\)的原根的充要条件是:

\[\forall 1 \leq j \leq s,g^{\frac{\varphi(m)}{q_j}} \neq 1 (\mod m)\]

证明:

根据欧拉定理,\(g^{\varphi(m)}=1(\bmod m)\)

根据定理3.1,\(\delta_m(g)|\varphi(m)\).如果\(\varphi(m)\)的每个非自身的因数\(d\)都不满足条件\(a^d=1(\bmod m)\)，那么根据原根的最小性,\(g\)就是原根。实际上我们不需要枚举每个因数,只需要枚举那些包含质因数最多的就可以了,因为如果有比\(\frac{\varphi(m)}{q_j}\)小的数满足条件,同样根据定理3.1,\(\frac{\varphi(m)}{q_j}\)也满足条件。也就是说,所有\(\frac{\varphi(m)}{q_j}\)的因数集合的并集包含了除\(\varphi(m)\)外\(\varphi(m)\)的所有因数。

定理4.3的推论设\(p\)是奇素数,\(p-1\)的所有的不同的素因数为\(q_1,q_2 \dots q_s\),那么\(g\)是模\(p\)的原根的充要条件是:

\[\forall 1 \leq j \leq s,g^{\frac{p-1}{q_j}} \neq 1 (\mod p)\]

求原根没有什么太好的算法，只能暴力枚举,然后根据定理4.3判断，但好在原根一般都很小,很快能找到答案。

下标给出了一些常用大质数的原根，容易发现原根都很小

p	原根g
3	2
5	2
17	3
97	5
193	5
257	3
7681	17
12289	11
40961	3
65537	3
786433	10
5767169	3
7340033	3
23068673	3
104857601	3
167772161	3
469762049	3
998244353	3
1004535809	3

NTT

NTT的定义

快速数论变换(Number Theory Transform,NTT)可以用于在模意义下计算两个整数序列的卷积，避免了FFT的精度问题。

从单位根到原根

我们发现,FFT用到了复数单位根的五大性质

\(\omega_n^0=\omega_n^n=1\)
\(\forall i \neq j,0<i,j<n,\omega_n^i \neq \omega_n^j\)
\(\omega_{2n}^{2k}=\omega_n^k\)
\(\omega_n^{k+\frac{n}{2}}=-\omega_n^k\)
\[\forall v \in \mathbb{N},\frac{1}{n} \sum_{k=0}^{n-1} \omega_n^{v k}=[v \bmod n=0]\]

想到我们之前提到的原根,如果把单位根换成原根,结果会如何呢?

一般的NTT的模数只能是形如\(p=a \times 2^b+1\)的质数.设\(p\)的原根为\(g\),我们定义\(g_n^k=(g^{\frac{p-1}{n}})^{k} (\bmod p)\)

下面我们证明\(g\)满足单位根的5条性质,注意所有计算在模\(p\)下进行:

定理5.1:\(g_n^0=g_n^n=1 (\bmod p)\)

证明:显然\(g_n^0=1\),\(g_n^n=g^{p-1}=g^{\varphi(p)}=1 (\bmod p)\),最后一步用到了原根的定义

定理5.2 \(\forall i \neq j(\bmod (p-1)),g_n^i \neq g_n^j (\bmod p)\)

证明:根据定理4.1,\(g^0,g^1 \dots g^{p-2}\)遍历\(p\)的剩余系,证毕。一定要注意,\(i,j\)是在模\(p-1\)的意义下

定理5.3 \(g_{2n}^{2k}=g_n^k\)

证明:\[g_{2 n}^{2 k} \equiv\left(g^{\frac{p-1}{2 n}}\right)^{2 k} \equiv\left(g^{\frac{p-1}{n}}\right)^{k} \equiv g_{n}^{k} \quad(\bmod p)\]

定理5.4:\(g_n^{k+\frac{n}{2}}=-g_n^k (\bmod p)\)

由于\(p\)是质数,且\(g_n^n=1 (\bmod p)\),因此\(g_n^{n/2}= ±1\)

由定理\(5.2\)得\(g_n^{n/2} \neq g_n^1=1\)

所以\(g_n^{n/2}=-1 (\bmod p)\),\(g_n^{k+\frac{n}{2}}=-g_n^k\)

定理5.5 \[\forall v \in \mathbb{N},\frac{1}{n} \sum_{k=0}^{n-1} \omega_n^{v k}=[v \bmod n=0]\]

当\(v \bmod n \neq 0\)时,根据等比数列求和公式

\[\begin{aligned} \sum_{k=0}^{n-1}\left(g_{n}^{v}\right)^{k} &=\frac{\left(g_{n}^{v}\right)^{n}-1}{g_{n}^{v}-1} \\ &=\frac{\left(g_{n}^{n}\right)^{v}-1}{g_{n}^{v}-1} \\ &=\frac{1-1}{g_{n}^{v}-1} \\ &=0 \end{aligned}\]

当\(v \mod n=0\)时,\(g_n^v=g^{\frac{(p-1)v}{n}}\),那么\((p-1)|\frac{(p-1)v}{n}\),根据定理3.1,\(g_n^v=1\),故\(g_n^{vk}=1\), \(\sum_{i=0}^{n-1}\left(g_{n}^{v}\right)^{i}=\sum_{i=0}^{n-1} 1=n\)

证毕。

至此，我们发现可以用\(g_n\)替代单位根. 注意到\(\frac{p-1}{n}\)必须是整数,考虑到FFT的过程是不断二分,那么\(p-1\)必须有足够的质因子2,这就是为什么一般的NTT模数都是\(p=a \times 2^b+1\)的形式

常用NTT模数表

\(p=a \times 2^b+1\)	\(a\)	\(b\)	原根\(g\)
3	1	1	2
5	1	2	2
17	1	4	3
97	3	5	5
193	3	6	5
257	1	8	3
7681	15	9	17
12289	3	12	11
40961	5	13	3
65537	1	16	3
786433	3	18	10
5767169	11	19	3
7340033	7	20	3
23068673	11	21	3
104857601	25	22	3
167772161	5	25	3
469762049	7	26	3
998244353	119	23	3
1004535809	479	21	3

节选自该链接

NTT的实现

void NTT(ll *x,int n,int type){
    for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]); 
    for(int len=1;len<n;len*=2){
        int sz=len*2;
        ll gn1=fast_pow((type==1?G:invG),(mod-1)/sz);
        for(int l=0;l<n;l+=sz){
            int r=l+len-1;
            ll gnk=1;
            for(int i=l;i<=r;i++){
                ll tmp=x[i+len];
                x[i+len]=(x[i]-gnk*tmp%mod+mod)%mod;
                x[i]=(x[i]+gnk*tmp%mod)%mod;
                gnk=gnk*gn1%mod; 
            }
        }
    } 
    if(type==-1){
        ll invn=inv(n);
        for(int i=0;i<n;i++) x[i]=x[i]*invn%mod; 
    }
}

可以与FFT代码对比

快速傅里叶变换(FFT)学习笔记(其二)(NTT)