基于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^{\frac{p-1}{p_i}}\neq1\pmod p$$

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

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^n=g^{nq}=g^{p-1}\equiv1\pmod p$$

又因为$(\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$：

$$S(\omega_n^k)=\sum_{i=0}^{n-1}(\omega_n^k)^i$$

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

$$S(\omega_n^k)=\frac{(\omega_n^k)^n-1}{\omega_n^k-1}$$

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

实现

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

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

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$

参考资料

• 从傅里叶变换到数论变换 - Menci
• 【数论】【原根】原根的性质以及如何求原根 - AutSky_JadeK
• 数学选讲.pptx - 成都七中培训day7
• 快速傅里叶变换.pptx - liuguangzhe