隐藏
基于FFT的NTT学习笔记 | Bill Yang's Blog

路终会有尽头,但视野总能看到更远的地方。

0%

基于FFT的NTT学习笔记

阅读本文章需要先学习快速傅里叶变换FFT

快速数论变换(Number Theoretic Transform,简称NTT),是一种使用模数原根为工具的离散傅里叶变换的算法。
我们注意到复数运算存在精度误差,因此对于只有整数参与的多项式运算,有时使用NTT是更好的选择。

FFT所用到单位复根性质

NTT基于FFT的过程,仅仅将单位复根工具替换为了数论原根,因此在替换的时候我们需要考虑FFT使用到了单位复根的什么性质。

  • $\omega_n^t(0\le t\le n-1)$互不相同,保证点值表示的合法。
  • 消去引理:对于任意的正整数$n,d$和非负整数$k$,满足$\omega_{dn}^{dk}=\omega_d^k$。
  • 折半引理:若$n$是偶数,有$\omega_n^{k+\frac n2}=-\omega_n^k$。
  • 求和引理:若对于正整数$n,k$,$n\nmid k$,有$\sum\limits_{j=0}^{n-1}(\omega_n^k)^j=0$

原根

定义

对于一个模数$p$,记其原根为$g$,满足$g^i\bmod p(0\le i\le n-1)$互不相同。

原根存在条件

并不是所有模数都存在原根,模$m$有原根的充要条件为:$m=2,4,p^a,2p^a$,其中$p$是奇素数,$a\ge1$。
根据条件,素数一定存在原根。

求原根的方法

对$p-1$唯一分解,写作$p-1=p_1^{a_1}p_2^{a_2}\cdots p_k^{a_k}$。
若对于任意$p_i$,有

恒成立,则$g$是$p$的原根。
若对于合数求原根,将$p-1$换为$\varphi(p)$。
因此我们可以在$O(\sqrt p)$的时间内求出所有原根。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
vector<int> ppt;

void find_root(LL x) {
LL tmp=x-1;
for(int i=2; i<=sqrt(tmp); i++)
if(tmp%i==0) {
ppt.push_back(i);
while(tmp%i==0)tmp/=i;
}
for(g=2; g<x; g++) {
bool bj=1;
for(int t:ppt)
if(Quick_Pow(g,(x-1)/t)==1) {
bj=0;
break;
}
if(bj)return;
}
}

原根的个数

一个数$m$若存在原根,则其原根个数为$\varphi(\varphi(m))$,特别地,对素数有$\varphi(p)=p-1$。
因此我们可以在$O(\sqrt n)$的时间内求出原根个数,在$O(n)$时间内求出$[1,n]$中任意数的原根个数。

NTT

我们尝试使用数论原根替代单位复根。
我们仍然使用$\omega_n$表示代入的$x$值。

互不相同

根据数论原根的定义,可以保证点值表示的合法性。

消去引理

由$\omega_n=g^q$可知$\omega_{2n}=g^\frac q2$($p=\frac q2\times 2n+1$),故$\omega_{2n}^{2k}=g^{2k\frac q2}=g^{kq}=\omega_n^k$,满足消去引理

折半引理

根据费马小定理得:

又因为$(\omega_n^\frac n2)^2=\omega_n^n$,所以$\omega_n^\frac n2\equiv\pm1\pmod p$,而根据性质一可得$\omega_n^\frac n2\neq\omega_n ^0$,即$\omega_n^\frac n2 \equiv -1\pmod p$。可推出$\omega_n^{k+\frac n2}=\omega_n^k\times\omega_n^\frac n2\equiv-\omega_n^k\pmod p$,满足折半引理

求和引理

当$k\neq0$时,$n\nmid k$:

根据等比数列求和公式得:

显然,$(\omega_n^k)^n-1\equiv0\pmod p$,故$S(\omega_n^k)=0\pmod p$,求和引理成立。

实现

至此,我们证明了数论原根满足所有FFT中用到的单位复根的性质,因此我们可以使用数论原根实现傅里叶变换。

具体的方案是,将单位根$\omega_n$替换为$g^q$,预处理$\omega_n^i$与其逆元,有乘法的时候取模即可。

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
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
const int maxn=131072+5;

LL Quick_Pow(LL a,LL b) {
LL sum=1;
for(; b; b>>=1,a=a*a%mod)if(b&1)sum=sum*a%mod;
return sum;
}

LL inv(LL x) {
return Quick_Pow(x,mod-2);
}

vector<int> ppt;

void find_root(LL x) {
LL tmp=x-1;
for(int i=2; i<=sqrt(tmp); i++)
if(tmp%i==0) {
ppt.push_back(i);
while(tmp%i==0)tmp/=i;
}
for(g=2; g<x; g++) {
bool bj=1;
for(int t:ppt)
if(Quick_Pow(g,(x-1)/t)==1) {
bj=0;
break;
}
if(bj)return;
}
}

struct NumberTheoreticTransform {
int n,rev[maxn];
LL omega[maxn],iomega[maxn];

void init(int n) {
this->n=n;
int x=Quick_Pow(g,(mod-1)/n);
omega[0]=iomega[0]=1;
for(int i=1; i<n; i++) {
omega[i]=omega[i-1]*x%mod;
iomega[i]=inv(omega[i]);
}
int k=log2(n);
for(int i=0; i<n; i++) { //二进制位翻转
int t=0;
for(int j=0; j<k; j++)if(i&(1<<j))t|=(1<<(k-j-1));
rev[i]=t;
}
}

void transform(LL* a,LL* omega) {
for(int i=0; i<n; i++)if(i<rev[i])swap(a[i],a[rev[i]]); //避免多次翻转
for(int len=2; len<=n; len*=2) {
int mid=len>>1;
for(LL* p=a; p!=a+n; p+=len)
for(int i=0; i<mid; i++) {
LL t=omega[n/len*i]*p[mid+i]%mod;
p[mid+i]=(p[i]-t+mod)%mod;
p[i]=(p[i]+t)%mod;
}
}
}

void dft(LL* a) {
transform(a,omega);
}

void idft(LL* a) {
transform(a,iomega);
LL x=inv(n);
for(int i=0; i<n; i++)a[i]=a[i]*x%mod;
}
} ntt;

void Multiply(LL* a1,const int n1,LL* a2,const int n2,LL* ans) {
int n=1;
while(n<n1+n2)n<<=1; //补成整数
ntt.init(n);
ntt.dft(a1);
ntt.dft(a2);
for(int i=0; i<n; i++)a1[i]=a1[i]*a2[i]%mod;
ntt.idft(a1);
for(int i=0; i<n1+n2-1; i++)ans[i]=a1[i];
}

FFT与NTT的选择

在讲这两者的选择,我们先来看看这两者的优缺点。

  • FFT的专长:实数多项式运算
  • FFT的不足:精度误差

  • NTT的专长:整数多项式运算

  • NTT的不足:不能参与实数多项式运算,对模数有要求

因此我们可以根据以下条件作出选择。

  • 题目给定模数$p$,$p-1$唯一分解后素数$2$的幂次较大,使用NTT。
  • 题目未指定模数,所有多项式运算只有整数参与,令模数$p=998244353$,使用NTT。
  • 否则使用FFT。

NTT常见模数与其最小原根

  • $p=998244353$,$g=3$
  • $p=1004535809$,$g=3$
  • $p=1005060097$,$g=5$

参考资料

姥爷们赏瓶冰阔落吧~