【NTT】任意模数NTT，两种方法的模板以及例题

最新推荐文章于 2023-07-10 19:39:24 发布

BUAA_Alchemist

最新推荐文章于 2023-07-10 19:39:24 发布

阅读量996

点赞数

分类专栏： FFT以及相关模板文章标签： NTT MTT

本文链接：https://blog.csdn.net/BUAA_Alchemist/article/details/86155114

版权

模板同时被 2 个专栏收录

7 篇文章 0 订阅

订阅专栏

FFT以及相关

3 篇文章 0 订阅

订阅专栏

模板题的地址：https://www.luogu.org/problemnew/show/P4245
给定 2 个多项式 F(x), G(x) ，请求出 F(x) * G(x)
系数对 p 取模，且不保证 p 可以分解成 p = a*2^k +1 之形式。
方法一：三模数法，这种方法相当于先求出真值，再进行取模运算，参见 https://blog.csdn.net/qq_35950004/article/details/79477797
（做了一些小修改：取模得当的话，可以不用快速乘法的）

#include<cstdio>
#include<algorithm>
using namespace std;
using LL=long long;

const int p[3]={998244353,1004535809,469762049};  //三个NTT模数
const int root[3]={3,3,3};  //三个原根

int n,m,a[3][1<<18],b[3][1<<18],mo,fa[100005],fb[100005];  //fa,fb是原多项式，a,b分别记录三个模数下DFT一次的答案
int bit,s,rev[1<<18];
LL invp0,invp1,invp01,c[1<<18];  //三个逆元是之后需要用到的，c是要求出的答案

int quick_power(int a, int b, int mo)
{
	int res=1,base=a;
	while(b)
	{
		if(b&1)
			res=(LL)res*base%mo;
		base=(LL)base*base%mo;
		b>>=1;
	}
	return res;
}

void get_rev(int bit)
{
	for(int i=0;i<(1<<bit);i++)
		rev[i]=(rev[i>>1]>>1)|((i&1)<<bit-1);
}

void FFT(int *a, int n, int dft, int o)
{
	int x,y;
	for(int i=0;i<n;i++)
		if(i<rev[i])
			swap(a[i],a[rev[i]]);
	for(int stp=1;stp<n;stp<<=1)
	{
		int wn=quick_power(root[o],(p[o]-1)/(stp*2),p[o]);
		if(dft==-1)
			wn=quick_power(wn,p[o]-2,p[o]);
		for(int j=0;j<n;j+=stp<<1)
		{
			int wnk=1;
			for(int k=j;k<j+stp;k++)
			{
				x=a[k];
				y=(LL)wnk*a[k+stp]%p[o];
				a[k]=(x+y)%p[o];
				a[k+stp]=(x-y+p[o])%p[o];
				wnk=(LL)wnk*wn%p[o];
			}
		}
	}
	if(dft==-1)
	{
		int t=quick_power(n,p[o]-2,p[o]);
		for(int i=0;i<n;i++)
			a[i]=(LL)a[i]*t%p[o];
	}
}

inline LL CRT01(int &ra, int &rb)  //合并p0和p1两个模数下的答案，用的是中国剩余定理
{
	return (invp1*ra%p[0]*p[1]+invp0*rb%p[1]*p[0])%((LL)p[0]*p[1]);
}

int main()
{
	scanf("%d%d%d",&n,&m,&mo);
	++n,++m;
	invp0=quick_power(p[0],p[1]-2,p[1]);
	invp1=quick_power(p[1],p[0]-2,p[0]);
	invp01=quick_power((LL)p[0]*p[1]%p[2],p[2]-2,p[2]);
	s=2; 
	for(bit=1;(1<<bit)<n+m-1;bit++)
		s<<=1;
	get_rev(bit);
	for(int i=0;i<n;i++)
		scanf("%d",&fa[i]);
	for(int i=0;i<m;i++)
		scanf("%d",&fb[i]);
	for(int i=0;i<3;i++)  //求fa的DFT
	{
		for(int j=0;j<n;j++)
			a[i][j]=fa[j];
		FFT(a[i],s,1,i);
	}
	for(int i=0;i<3;i++)  //求fb的DFT
	{
		for(int j=0;j<m;j++)
			b[i][j]=fb[j];
		FFT(b[i],s,1,i);
	}
	for(int i=0;i<3;i++)  //卷积之后IDFT，求出三个模数下的答案
	{
		for(int j=0;j<s;j++)
			a[i][j]=(LL)a[i][j]*b[i][j]%p[i];
		FFT(a[i],s,-1,i);
	}
	for(int i=0;i<n+m-1;i++)  //合并前两个模数下的答案
		c[i]=CRT01(a[0][i],a[1][i]);
	LL K;
	for(int i=0;i<n+m-1;i++)  //求出满足K*p1*p0 % p2=a2-a1等式的K，可以证明这时K只能取小于p2的非负数
	{
		K=(a[2][i]-c[i]%p[2]+p[2])%p[2];	
		c[i]=((K*invp01%p[2]*p[0]%mo*p[1]%mo+c[i])%mo+mo)%mo;
	}
	for(int i=0;i<n+m-1;i++)
		printf("%lld ",c[i]);
	return 0;
}

方法二：
处理任意模数NTT问题
令M=sqrt(mod)
然后多项式的每一项拆成AM+B，于是A和B都在int之内就不那么容易爆精度了
所以两个数相乘就成为了
(A1M+B1)∗(A2M+B2)=A1A2M2+(A1B2+A2B1)M+B1B2
分别进行4次DFT和3次IDFT即可
不得不说这种方法有点冒险，因为实际上精度还是不好保证，这里不仅用了long double，而且事先要对系数取模，过了有点侥幸。这种比较好些，不过要稳的话还是采用第一种方法
代码比较易懂，就不加注释了

#include<cstdio>
#include<complex>
#include<cmath>
using namespace std;
using db=long double;
using LL=long long;
using cp=complex<db>;

const db PI=acos(-1.0L);
 
int rev[1<<18],s,Base,n,m,mo,bit;
cp aA[1<<18],aB[1<<18],bA[1<<18],bB[1<<18],p[1<<18],q[1<<18],r[1<<18];

void get_rev(int bit)
{
	for(int i=0;i<(1<<bit);i++)
		rev[i]=(rev[i>>1]>>1)|((i&1)<<bit-1);
}

void FFT(cp *a, int n, int dft)
{
	cp x,y;
	for(int i=0;i<n;i++)
		if(i<rev[i])
			swap(a[i],a[rev[i]]);
	for(int stp=1;stp<n;stp<<=1)
	{
		cp wn=exp(cp(0,dft*PI/stp));
		for(int j=0;j<n;j+=stp<<1)
		{
			cp wnk(1,0);
			for(int k=j;k<j+stp;k++)
			{
				x=a[k];
				y=wnk*a[k+stp];
				a[k]=x+y;
				a[k+stp]=x-y;
				wnk*=wn;
			}
		}
	}
	if(dft==-1)
		for(int i=0;i<n;i++)
			a[i]/=n;
}

int main()
{
	scanf("%d%d%d",&n,&m,&mo);
	++n,++m;
	s=2;
	Base=sqrt(mo)+1;
    for(bit=1;(1<<bit)<n+m-1;bit++)
        s<<=1;
    get_rev(bit);
    for(int i=0,fa;i<n;i++)
        scanf("%d",&fa),fa%=mo,aA[i]=(db)(fa/Base),aB[i]=(db)(fa%Base);
    for(int i=0,fb;i<m;i++)
        scanf("%d",&fb),fb%=mo,bA[i]=(db)(fb/Base),bB[i]=(db)(fb%Base);
	FFT(aA,s,1);
	FFT(aB,s,1);
	FFT(bA,s,1);
	FFT(bB,s,1);
	for(int i=0;i<s;i++)
	{
		p[i]=aA[i]*bA[i];
		q[i]=aA[i]*bB[i]+aB[i]*bA[i];
		r[i]=aB[i]*bB[i];
	}
	FFT(p,s,-1);
	FFT(q,s,-1);
	FFT(r,s,-1);
	for(int i=0,ans;i<n+m-1;i++)
	{
		ans=(LL)(p[i].real()+0.5)*Base%mo*Base%mo;
		ans=(ans+(LL)(q[i].real()+0.5)*Base%mo)%mo;
		ans=(ans+(LL)(r[i].real()+0.5))%mo;
		printf("%d ",ans);
	}
    return 0;
}

补充一道非常好的题目
https://www.luogu.org/problemnew/show/P5173

题目背景
临近中考，pG的班主任决定上一节体育课，放松一下。

题解：https://blog.csdn.net/kkkksc03/article/details/85008120

题目描述
老师带着pG的同学们一起做传球游戏。

游戏规则是这样的： nn 个同学站成一个圆圈，其中的一个同学手里拿着一个球，当老师吹哨子时开始传球，每个同学可以把球传给自己左右的两个同学中的一个（左右任意），当老师再次吹哨子时，传球停止，此时，拿着球没有传出去的那个同学就是败者，要给大家表演一个节目。

pG提出一个有趣的问题：有多少种不同的传球方法可以使得从pG手里开始传的球，传了 mm 次以后，又回到pG手里。两种传球方法被视作不同的方法，当且仅当这两种方法中，接到球的同学按接球顺序组成的序列是不同的。比如有三个同学 11 号、 22 号、 33 号，并假设pG为 11 号，球传了 33 次回到pG手里的方式有 1 -> 2 -> 3 -> 11−>2−>3−>1 和 1 -> 3 -> 2 -> 11−>3−>2−>1 ，共22 种。

输入输出格式
输入格式：
一行，有两个用空格隔开的整数 n,mn,m
输出格式：
11 个整数，表示符合题意的方法数。

由于答案可能过大，对10^9+710
9
+7取模。

输入输出样例
输入样例#1：
3 3
输出样例#1：
2
输入样例#2：
30 30
输出样例#2：
155117522
输入样例#3：
1234 12345678
输出样例#3：
424074635
说明
对于8%的数据，n \le 100,m \le 10^4n≤100,m≤10
4
.

对于100%的数据，n \le 3500,m \le 10^9n≤3500,m≤10
9
.

首先很容易想到杨辉三角和组合数，然而需要累和的组合数太多了，这么做会TLE，其次很容易想到矩阵乘法或者FFT优化DP。矩阵乘法仍然有T的风险，这里选用FFT，因为答案数据很大，只能用任意模数的NTT来做。这里方便起见，我们把递推的多项式定义成01000000……1，这样，累次乘法之后首个位置就是答案了。
用类似快速幂的方法做logN次任意模数NTT。需要注意的是，因为这里的乘法是“循环的”，每乘一次还要整理一遍，把后面的n——s-1个数字叠加在前面0——n-1上。
另外，不能够把整合三个答案，对1E9+7取模这些步骤放到最后进行（本来想偷懒的），必须每次乘完就要整合出当下的答案。这是因为三模数NTT的正确性是建立于于真实答案的范围不超过三模数乘积的前提下的，如果一次乘法做完后不及时对1E9+7取模，真实的值就会远远超出三模数乘积，大数据下对1E9+7取模的答案会与真实答案有偏差。
效率大概是O(nlognlogm)级别的。
附上代码，另外，不知道哪里写臭了，这份代码表现欠佳，差点T了，以后再作改进。

#include<cstdio>
#include<algorithm>
using namespace std;
using LL=long long;

const int p[3]={998244353,1004535809,469762049}; 
const int root[3]={3,3,3};
const int mo=1000000007;

int n,m,a[3][1<<13];
int bit,s,rev[1<<13],base[3][1<<13];
LL invp0,invp1,invp01,c[3505];

int quick_power(int a, int b, int mo)
{
    int res=1,base=a;
    while(b)
    {
        if(b&1)
            res=(LL)res*base%mo;
        base=(LL)base*base%mo;
        b>>=1;
    }
    return res;
}

void get_rev(int bit)
{
    for(int i=0;i<(1<<bit);i++)
        rev[i]=(rev[i>>1]>>1)|((i&1)<<bit-1);
}

void FFT(int *a, int n, int dft, int o)
{
    int x,y;
    for(int i=0;i<n;i++)
        if(i<rev[i])
            swap(a[i],a[rev[i]]);
    for(int stp=1;stp<n;stp<<=1)
    {
        int wn=quick_power(root[o],(p[o]-1)/(stp*2),p[o]);
        if(dft==-1)
            wn=quick_power(wn,p[o]-2,p[o]);
        for(int j=0;j<n;j+=stp<<1)
        {
            int wnk=1;
            for(int k=j;k<j+stp;k++)
            {
                x=a[k];
                y=(LL)wnk*a[k+stp]%p[o];
                a[k]=(x+y)%p[o];
                a[k+stp]=(x-y+p[o])%p[o];
                wnk=(LL)wnk*wn%p[o];
            }
        }
    }
    if(dft==-1)
    {
        int t=quick_power(n,p[o]-2,p[o]);
        for(int i=0;i<n;i++)
            a[i]=(LL)a[i]*t%p[o];
    }
}

inline LL CRT01(int &ra, int &rb)
{
    return (invp1*ra%p[0]*p[1]+invp0*rb%p[1]*p[0])%((LL)p[0]*p[1]);
}

void mul(int *a, int *b, int n, int s, int o)
{
    FFT(a,s,1,o);
    if(a!=b)
        FFT(b,s,1,o);
    for(int i=0;i<s;i++)
        a[i]=(LL)a[i]*b[i]%p[o];
    FFT(a,s,-1,o);
    for(int i=s-1;i>=n;i--)
        a[i-n]=(a[i-n]+a[i])%p[o],a[i]=0;
    if(a!=b)
        FFT(b,s,-1,o);
}

void solve(int a[][1<<13], int b[][1<<13], int n, int s)
{
	for(int i=0;i<3;i++)
        mul(a[i],b[i],n,s,i);
    for(int i=0;i<n;i++)    //每次mul完，都要及时整合出临时答案
	    c[i]=CRT01(a[0][i],a[1][i]);
	LL K;
    for(int i=0;i<n;i++)
    {
        K=(a[2][i]-c[i]%p[2]+p[2])%p[2];	
        c[i]=(K*invp01%p[2]*p[0]%mo*p[1]%mo+c[i])%mo;
        for(int j=0;j<3;j++)
	        a[j][i]=c[i];
 	}
}

int main()
{
    scanf("%d%d",&n,&m);
    invp0=quick_power(p[0],p[1]-2,p[1]);
    invp1=quick_power(p[1],p[0]-2,p[0]);
    invp01=quick_power((LL)p[0]*p[1]%p[2],p[2]-2,p[2]);
    s=2;
    for(bit=1;(1<<bit)<n+n-1;bit++)
        s<<=1;
    get_rev(bit);
    for(int i=0;i<3;i++)
    	a[i][0]=base[i][1]=base[i][n-1]=1;
    while(m>0)
    {
        if(m&1)
        	solve(a,base,n,s);
        solve(base,base,n,s);
        m>>=1;
    }
    printf("%lld ",a[0][0]);
    return 0;
}

BUAA_Alchemist

关注

0
点赞
踩
2

收藏

觉得还不错? 一键收藏
0
评论
【NTT】任意模数NTT，两种方法的模板以及例题

模板题的地址：https://www.luogu.org/problemnew/show/P4245给定 2 个多项式 F(x), G(x) ，请求出 F(x) * G(x)系数对 p 取模，且不保证 p 可以分解成 p = a*2^k +1 之形式。方法一：三模数法，参见 https://blog.csdn.net/qq_35950004/article/details/79477797...
复制链接

扫一扫

专栏目录