[学习笔记] 快速数论变换 (NTT)

目录

• 0. 前置芝士
• 1. 正文
  • 1.1. 转化
  • 1.2. 质数的选取
  • 1.3. 代码实现

0. 前置芝士

可以看看我写的 \(\mathtt{FFT}\)​ 还有 原根。

1. 正文

1.1. 转化

还记得单位复根吗？事实上，如果我们找到一个数满足单位复根的性质，就可以将其替代。

定义 \(\omega_n\) 为 \(g^{\frac{\varphi(p)}{n}}\)​.

1. \(\omega_n^j=\omega_{n\times k}^{j\times k}\). 这就是 \(g\) 的指数分子分母同时乘上 \(k\)，所以相等。
2. \(\omega_n^j=-\omega_n^{j+\frac{n}{2}}\). \(\omega_n^{\frac{n}{2}}\) 就是 \(g^{\frac{\varphi(p)}{2}}\) 即 \(-1\)​.
3. \(\omega_n^n=1\). 显然 \(g^{\varphi(p)}=1\)​.

1.2. 质数的选取

由于单位根是 \(g^{\frac{\varphi(p)}{n}}\)，所以要保证 \(\varphi(p)\) 含有足够的 \(2\) 的幂。

1.3. 代码实现

注意 \(\mathtt{P3803}\)​ 是没有模数的，不过数据范围太小，可以看作是取模之后的结果。

#include <cstdio>
#define print(x,y) write(x),putchar(y)

template <class T>
inline T read(const T sample) {
	T x=0; char s; bool f=0;
	while((s=getchar())>'9' || s<'0')
		f |= (s=='-');
	while(s>='0' && s<='9')
		x = (x<<1)+(x<<3)+(s^48),
		s = getchar();
	return f?-x:x;
}

template <class T>
inline void write(T x) {
	static int writ[50],tp=0;
	if(x<0) putchar('-'),x=-x;
	do writ[++tp] = x-x/10*10, x/=10; while(x);
	while(tp) putchar(writ[tp--]^48);
}

#include <iostream>
using namespace std;

const int maxn = 4e6+5;
const int mod = 998244353, G = 3;

int n,m,lim=1,bit,wn,w,tmp;
int a[maxn],b[maxn],rev[maxn],iG;

inline int qkpow(int x,int y=mod-2) {
	int r=1;
	while(y) {
		if(y&1) r = 1ll*r*x%mod;
		x = 1ll*x*x%mod; y>>=1;
	}
	return r;
}

void preWork() {
	iG = qkpow(G);
	while(lim<n+m-1) lim<<=1,++bit;
	for(int i=0;i<lim;++i)
		rev[i] = (rev[i>>1]>>1)|((i&1)<<bit-1);
}

inline int inc(int x,int y) {
	return x+y>=mod?x+y-mod:x+y;
}

inline int dec(int x,int y) {
	return x-y<0?x-y+mod:x-y;
}

void NTT(int *f,bool opt) {
	for(int i=0;i<lim;++i)
		if(i<rev[i]) swap(f[i],f[rev[i]]);
	for(int mid=1;mid<lim;mid<<=1) {
		wn = qkpow(opt?G:iG,(mod-1)/(mid<<1));
		for(int i=0;i<lim;i+=(mid<<1)) {
			w = 1;
			for(int j=0;j<mid;++j,w=1ll*w*wn%mod) {
				tmp = 1ll*f[i+j+mid]*w%mod;
				f[i+j+mid] = dec(f[i+j],tmp);
				f[i+j] = inc(f[i+j],tmp);
			}
		}
	}
}

int main() {
	n=read(9)+1,m=read(9)+1;
	for(int i=0;i<n;++i)
		a[i] = read(9);
	for(int i=0;i<m;++i)
		b[i] = read(9);
	preWork();
	NTT(a,1),NTT(b,1);
	for(int i=0;i<lim;++i)
		a[i] = 1ll*a[i]*b[i]%mod;
	NTT(a,0);
	int Inv = qkpow(lim);
	for(int i=0;i<n+m-1;++i)
		printf("%d ",1ll*a[i]*Inv%mod);
	return 0;
}