洛谷 P4245 【模板】任意模數NTT

• 洛谷 P4245 【模板】任意模數NTT

• 貼個板子，4次DFT。

Code

#include<cstdio>
#include<algorithm>
#include<cmath>
#include<cctype>
using namespace std;
typedef long long LL;
const int N=1e5+5;
const double Pi=acos(-1);
struct comp
{
	double r,i;
	inline comp(){}
	inline comp(double rr,double ii){r=rr,i=ii;}
	friend inline comp operator +(comp x,comp y)
	{
		return comp(x.r+y.r,x.i+y.i);
	}
	friend inline comp operator -(comp x,comp y)
	{
		return comp(x.r-y.r,x.i-y.i);
	}
	friend inline comp operator *(comp x,comp y)
	{
		return comp(x.r*y.r-x.i*y.i,x.r*y.i+x.i*y.r);
	}
	friend inline comp operator /(comp x,double y)
	{
		return comp(x.r/y,x.i/y);
	}
	friend inline comp conj(comp x)
	{
		return comp(x.r,-x.i);
	}
};
int rev[N<<2],a[N],b[N],ans[N<<1];
comp a1[N<<2],b1[N<<2],a2[N<<2],b2[N<<2],wn[N<<2];
inline int read()
{
    int X=0,w=0; char ch=0;
    while(!isdigit(ch)) w|=ch=='-',ch=getchar();
    while(isdigit(ch)) X=(X<<3)+(X<<1)+(ch^48),ch=getchar();
    return w?-X:X;
}
void write(int x)
{
	if(x>9) write(x/10);
	putchar(x%10+'0');
}
inline void FFT(comp *y,int len,int ff)
{
	for(int i=0;i<len;i++)
		if(i<rev[i]) swap(y[i],y[rev[i]]);
	for(int h=2,d=len>>1;h<=len;h<<=1,d>>=1)
		for(int i=0,e=h>>1;i<len;i+=h)
			for(int j=0,cnt=0;j<e;j++,cnt+=d)
			{
				comp u=y[i+j],v=wn[cnt]*y[i+j+e];
				y[i+j]=u+v;
				y[i+j+e]=u-v;
			}
	if(ff==-1)
	{
		reverse(y+1,y+len);
		for(int i=0;i<len;i++) y[i]=y[i]/len;
	}
}
int main()
{
	int n=read()+1,m=read()+1,p=read();
	for(int i=0;i<n;i++) a[i]=read();
	for(int i=0;i<m;i++) b[i]=read();
	int len=1,ll=0;
	while(len<n+m) len<<=1,ll++;
	for(int i=0;i<len;i++) rev[i]=rev[i>>1]>>1|(i&1)<<ll-1;
	for(int i=0;i<=len;i++) wn[i]=comp(cos(2*Pi*i/len),sin(2*Pi*i/len));
	for(int i=0;i<n;i++) a1[i]=comp(a[i]>>15,a[i]&32767);
	for(int i=0;i<m;i++) b1[i]=comp(b[i]>>15,b[i]&32767);
	FFT(a1,len,1),FFT(b1,len,1);
	for(int i=0;i<len;i++)
	{
		int j=len-i&len-1;
		comp ea1=(a1[i]+conj(a1[j]))*comp(0.5,0);
		comp ea0=(a1[i]-conj(a1[j]))*comp(0,-0.5);
		comp eb1=(b1[i]+conj(b1[j]))*comp(0.5,0);
		comp eb0=(b1[i]-conj(b1[j]))*comp(0,-0.5);
		a2[i]=ea1*eb1+ea1*eb0*comp(0,1);
		b2[i]=ea0*eb1+ea0*eb0*comp(0,1);
	}
	FFT(a2,len,-1),FFT(b2,len,-1);
	for(int i=0;i<n+m-1;i++)
	{
		int e1=(LL)round(a2[i].r)%p;
		int e2=(LL)round(a2[i].i)%p;
		int e3=(LL)round(b2[i].r)%p;
		int e4=(LL)round(b2[i].i)%p;
		ans[i]=((((LL)e1<<30)+((LL)e2+e3<<15)+e4)%p+p)%p;
	}
	for(int i=0;i<n+m-1;i++) write(ans[i]),putchar(' ');
	return 0;
}