CCPC-Wannafly Summer Camp Day 2 I Steins;Gate [原根+FFT]
阿新 • • 發佈:2018-12-13
最近牛客重開了wannafly camp的題,打算把之前不會的題目補掉
這個 要求 ai*aj %P =ak的種類數,我們對 p求原根,即為G 那麼 ai%p就可以用 % p得到
這樣這個表示式 就變成了 那麼我們只要求 在給定資料中 x+y =z的種類數有多少就行
這就是一個很明顯的FFT
然後零的時候單獨處理
#include <bits/stdc++.h> using namespace std; int N,P; typedef long long ll; const int maxn =800000+50; const long double pi=acos(-1.0); struct Complex{ long double x,y; Complex(long double _x=0,long double _y=0) { x=_x; y=_y; } Complex operator +(const Complex &b)const { return Complex(x+b.x,y+b.y); } Complex operator -(const Complex &b)const { return Complex(x-b.x,y-b.y); } Complex operator *(const Complex &b)const { return Complex(x*b.x-y*b.y,x*b.y+y*b.x);; } }a[maxn],b[maxn]; ll S,T,n,m,L,R[maxn],ans[maxn]; long long F[maxn]; void FFT(Complex a[],int opt) { for(int i=0;i<n;i++) if(i<R[i]) swap(a[i],a[R[i]]); for(int k=1;k<n;k<<=1) { Complex wn = Complex( cos(pi/k),opt*sin(pi/k) ); for(int i=0;i<n;i+=(k<<1)) { Complex w = Complex(1,0); for(int j=0;j<k;++j,w=w*wn) { Complex x = a[i+j], y = w*a[i+j+k]; a[i+j] = x+y; a[i+j+k] =x-y; } } } } void calc(int opt) { FFT(a,1); for(int i=0;i<=n;i++) a[i] = a[i]*a[i]; FFT(a,-1); for(int i=0;i<=n;i++) { F[i] = (long long)(a[i].x/n+0.5)*opt; } } bool is_root(int r,int p) { int x=1; for(int i=1;i<p-1;i++) { x=x*r%p; if(x==1) return 0; } return 1; } long long x[maxn]; long long idx[maxn],vis[maxn]; long long cnt[maxn]; int main() { scanf("%d%d",&N,&P); memset(cnt,0,sizeof(cnt)); memset(idx,0,sizeof(idx)); memset(vis,0,sizeof(vis)); memset(x,0,sizeof(x)); int root=0; for(int i=2;i<P;i++) { if(is_root(i,P)) { root = i; break; } } idx[0]=1; //cout<<root<<endl; for(int i=1;i<=P-1;i++) { idx[i]=idx[i-1]*root%P; vis[idx[i]]=i; } long long zcnt=0; for(int i=1;i<=N;i++) { scanf("%lld",&x[i]); if(x[i]%P==0) zcnt++; else cnt[vis[x[i]%P]]++; } m =2*P-2;L=0; for(n =1;n <=m;n<<=1) ++L; for(int i=0;i<n;i++) { R[i] = (R[i>>1]>>1)|((i&1)<<(L-1)); } for(int i=0;i<=n;i++) { a[i] = Complex(1.0*cnt[i],0.0); } calc(1); ll rs=1; for(int i=0;i<=m;i++) { //cout<<i<<" "<<F[i]<<" "<<endl; ans[vis[rs]]+=F[i]; //cout<<rs<<endl; rs = rs*root%P; } //cout<<N<<endl; ll allz = (ll)N*N-(ll)(N-zcnt)*(N-zcnt); for(int i=1;i<=N;i++) { if(x[i]>=P) printf("%d\n",0); else if(x[i]%P==0) printf("%lld\n",allz); else printf("%lld\n",ans[vis[x[i]]]); } return 0; }