CCPC-Wannafly Summer Camp Day 2 I Steins;Gate [原根+FFT]

最近牛客重开了wannafly camp的题，打算把之前不会的题目补掉

这个要求 ai*aj %P =ak的种类数，我们对 p求原根，即为G 那么 ai%p就可以用 $G^{t}$ % p得到

这样这个表达式就变成了 $G^{x} * G^{y} = G^{z}$ 那么我们只要求在给定数据中 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;
}