快速数论变换(NTT)

快速数论变换(NTT)

这东西之前就想学了,一直没有动手 orz,现在补一下。

学这东西我感觉并没有很多新知识,学之前掌握 FFT 就好了。

FFT 可以在这里看看:https://www.cnblogs.com/Tenshi/p/15434004.html

NTT,是用来解决多项式乘法取模问题的,因为 FFT 可能在精度上不够,而且常数较大,因此使用数论的手段将 FFT 进行改造得到 NTT。

原理

我们先假设模数 \(P\)​ 具有原根 \(rt\)​,这意味着 \(rt^{\varphi(P)}=1\)​,如果模数更特殊一点,假设模数为具有原根的质数(比如 \(998244353\))。那么则有 \(rt^{P-1} \equiv 1 \pmod P,~ rt^{P} \equiv rt \pmod P\)​​。

而在 FFT 中,我们有 \((\cos\theta + i\sin\theta)^{2\pi/ \theta} = 1\),形式上和 \(rt^{P-1} \equiv 1 \pmod P\) 非常接近,我们再对模数进行限制:模数需要能表述为 \(k\times 2^x + 1\)​ 的形式。(因为这样才能保证在分治的过程中 \(P-1\) 能被 \(2\) 的幂次整除)。

因此在取模意义下我们可以类似于 FFT 写出 NTT 的代码:

void NTT(ll *a, int type, int mod){
	for(int i=0; i<tot; i++){
		a[i]%=mod;
		if(i<rev[i]) swap(a[i], a[rev[i]]);
	}
	
	for(int mid=1; mid<tot; mid<<=1){
		ll w1=fpow(rt, (type==1? (mod-1)/(mid<<1): mod-1-(mod-1)/(mid<<1)), mod);
		for(int i=0; i<tot; i+=mid*2){
			ll wk=1;
			for(int j=0; j<mid; j++, wk=wk*w1%mod){
				auto x=a[i+j], y=wk*a[i+j+mid]%mod;
				a[i+j]=(x+y)%mod, a[i+j+mid]=(x-y+mod)%mod;
			}
		}
	}
	
	if(type==-1){
		for(int i=0; i<tot; i++) a[i]=a[i]*inv(tot, mod)%mod;
	}
}

拓展

由上所述,我们在 NTT 过程中对模数进行了一些限制,自然不能够直接推广来解决任意模数 NTT 的问题,怎么办呢?

结合模板题来说:

https://www.luogu.com.cn/problem/P4245

记题目所给的模数为 \(P\)。

我们可以选择三个模数 \(m_1,m_2,m_3\)​ 满足 \(m_1m_2m_3 > nP^2\)​。

可以取为 \(m_1=998244353,~ m_2=1004535809,~ m_3=469762049\)。​

我们先用这三个模数分别做一次 NTT,然后用 CRT(中国剩余定理)将这三个结果合并即可。

直接合并会爆 long long,那怎么合并呢?具体来说:

假设所求结果为 \(ans\),同时我们约定 \(inv(x,y)\) 代表 \(x\) 模 \(y\)​ 的逆元。

\[ans\equiv c_1 \pmod {m_1} \\ ans\equiv c_2 \pmod {m_2} \\ ans\equiv c_3 \pmod {m_3} \\ \]

先对前两个进行合并:

\(ans \equiv c_1\times m_2\times inv(m_1, m_2) + c_2\times m_1\times inv(m_2, m_1) \pmod {m_1m_2}\)。

记 \(M=m_1m_2, ~ C=c_1\times m_2\times inv(m_1, m_2) + c_2\times m_1\times inv(m_2, m_1)\)

那么我们设 \(ans = xM + C = ym_3 + c_3\)​​。

则有 \(x \equiv (c_3-C)\times inv(M, m_3) \pmod {m_3}\)​。

设 \(t = (c_3-C)\times inv(M, m_3)\)​​,\(t\) 显然可以直接算出来,我们进一步设 \(x = km_3 + t\)。

则 \(ans = km_1m_2m_3 + tM + C\)​​。

因为 \(ans < m_1 m_2 m_3\),因此 \(ans=tM+C\)。

最后对 \(P\) 取模即可。​

实现:

// Problem: P4245 【模板】任意模数多项式乘法
// Contest: Luogu
// URL: https://www.luogu.com.cn/problem/P4245
// Memory Limit: 500 MB
// Time Limit: 2000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

#include<bits/stdc++.h>
using namespace std;

#define debug(x) cerr << #x << ": " << (x) << endl
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define dwn(i,a,b) for(int i=(a);i>=(b);i--)

using pii = pair<int, int>;
using ll = long long;

inline void read(int &x){
    int s=0; x=1;
    char ch=getchar();
    while(ch<'0' || ch>'9') {if(ch=='-')x=-1;ch=getchar();}
    while(ch>='0' && ch<='9') s=(s<<3)+(s<<1)+ch-'0',ch=getchar();
    x*=s;
}

const int N=3e5+5;
const ll m1=998244353, m2=1004535809, m3=469762049, M=m1*m2, rt=3;

int n, m, P;
ll a[3][N], b[3][N], ans[N];

int rev[N], tot=1, bit;

ll fpow(ll x, int p, ll mod){
	int res=1;
	for(; p; p>>=1, x=x*x%mod) if(p&1) res=res*x%mod;
	return res;
}

ll inv(ll x, ll mod){
	return fpow(x, mod-2, mod);
}

ll mul(ll x, int p, ll mod){
	ll res=0;
	for(; p; p>>=1, x=(x+x)%mod) if(p&1) res=(res+x)%mod;
	return res;
}

void NTT(ll *a, int type, int mod){
	for(int i=0; i<tot; i++){
		a[i]%=mod;
		if(i<rev[i]) swap(a[i], a[rev[i]]);
	}
	
	for(int mid=1; mid<tot; mid<<=1){
		ll w1=fpow(rt, (type==1? (mod-1)/(mid<<1): mod-1-(mod-1)/(mid<<1)), mod);
		for(int i=0; i<tot; i+=mid*2){
			ll wk=1;
			for(int j=0; j<mid; j++, wk=wk*w1%mod){
				auto x=a[i+j], y=wk*a[i+j+mid]%mod;
				a[i+j]=(x+y)%mod, a[i+j+mid]=(x-y+mod)%mod;
			}
		}
	}
	
	if(type==-1){
		for(int i=0; i<tot; i++) a[i]=a[i]*inv(tot, mod)%mod;
	}
}

void CRT(){
	for(int i=0; i<tot; i++){
		ll res=0;
		(res+=mul(a[0][i]*m2%M, inv(m2, m1), M))%=M;
		(res+=mul(a[1][i]*m1%M, inv(m1, m2), M))%=M;
		a[1][i]=res;
	}
	for(int i=0; i<tot; i++){
		ll res=(a[2][i]-a[1][i]%m3+m3)%m3*inv(M%m3, m3)%m3;
		ans[i]=(M%P*res%P+a[1][i]%P)%P;
	}
}

void solve(int k, int mod){
	NTT(a[k], 1, mod), NTT(b[k], 1, mod);
	for(int i=0; i<tot; i++) a[k][i]=a[k][i]*b[k][i]%mod;
	NTT(a[k], -1, mod);
}

int main(){
	cin>>n>>m>>P;
	rep(i,0,n){
		int t; read(t);
		rep(j,0,2) a[j][i]=t%P;
	}
	rep(i,0,m){
		int t; read(t);
		rep(j,0,2) b[j][i]=t%P;
	}
	
	while(tot<=n+m) bit++, tot<<=1;
	for(int i=0; i<tot; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(bit-1));
	
	solve(0, m1), solve(1, m2), solve(2, m3);
	CRT();
	
	rep(i,0,n+m) cout<<ans[i]<<' ';
	cout<<endl;
	
	return 0;
}
上一篇:机器学习数学基础|Datawhale-8月(2)笔记


下一篇:Java使用JsonSchema实例,基于Springboot