1. 程式人生 > >牛客Steins;Gate(原根+FFT)

牛客Steins;Gate(原根+FFT)

9516: Steins;Gate

時間限制: 2 Sec  記憶體限制: 128 MB
提交: 33  解決: 10
[提交] [狀態] [討論版] [命題人:admin]

題目描述

助手作為物理學家,小時候當然參加過數學競賽(MO)啦。在助手還是小蘿莉的時候,她的數學老師曾經給她出了這麼一道題:
現有NN個數a1,a2,⋯,aN。對於每一個ak,求有多少個有序二元組(i,j)滿足(ai×aj) mod P=ak,其中P為一給定質數。

輸入

第一行有兩個正整數N,P(1≤N≤2×105,2≤P≤2×105),P為質數。
第二行N個非負整數a1,a2,⋯,aN(0≤ai≤2.1×109)。

輸出

按照格式輸出N個數,分別為a1,a2,⋯,aN對應的答案。

樣例輸入

7 3
0 2 0 1 2 3 4

樣例輸出

33
8
33
8
8
0
0

來源/分類

【分析】

p是質數,故一定存在原根g,且[1,p-1]之間的數x都可以表示為x=g^{i}。將每個數出現的次數 c[i] 作為係數,構成多項式f(g)=c_{0}g^{0}+c_{1}g^{1}+...+c_{n}g^{n},令F(g)=f(g)*f(g),則f(g)自我相乘之後,F(g)的各項係數,就是每個數字出現的次數。

根據題意計算F(g)的各項係數。多項式相乘用FFT加速

【程式碼】

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAX=1e6+5;
const double PI=acos(-1.0);
typedef complex<double> cpx;
ll qpow(ll n,ll m,ll p)
{
    ll ans=1;
    n%=p;
    for(;m;m>>=1,n=n*n%p)
        if(m&1)ans=ans*n%p;
    return ans;
}
int pri[MAX],tot;
ll getRoot(ll p) //求質數p的最小原根
{
    tot=0;
    ll n=p-1,sq=sqrt(p+0.5);
    for(int i=2;i<=sq;i++)if(n%i==0)
    {
        pri[tot++]=i;
        while(n%i==0)n/=i;
    }
    if(n>1)pri[tot++]=n;
    for(int g=2;g<=p-1;g++) //試探每一個g是否原根
    {
        int flag=1;
        for(int i=0;i<tot;i++)if(qpow(g,(p-1)/pri[i],p)==1)
        {
            flag=0; break;
        }
        if(flag)return g;
    }
    return -1; //沒有原根
}
 
int pos[MAX]; //DFT變換位置。結果一定是兩者互換,不存在三者互換
void initpos(int n)
{
    int k=log2(n+0.5)-1;
    for(int i=0;i<n;i++) //位置變換
        pos[i]=(pos[i>>1]>>1)|(i&1)<<k;
}
void FFT(cpx a[],int n,int tag) //多項式項數n,請確保n是2的次冪
{
    initpos(n); //一定要預處理pos
    for(int i=0;i<n;i++)if(i<pos[i])
        swap(a[i],a[pos[i]]);
    for(int k=2;k<=n;k<<=1) //列舉分治合併過程塊長度
    {
        int m=k>>1; //半個塊的大小
        cpx wn(cos(tag*2*PI/k),sin(tag*2*PI/k)); //單位k次復根
        for(int i=0;i<=n;i+=k) //列舉分治塊
        {
            cpx now(1,0); //單位復根的0次冪
            for(int j=0;j<m;j++) //分治塊內的每一項
            {
                cpx temp=now*a[i+j+m];
                a[i+j+m]=a[i+j]-temp;
                a[i+j]=a[i+j]+temp;
                now=now*wn;  //旋轉
            }
        }
    }
    if(tag==-1)
        for(int i=0;i<n;i++)a[i]=a[i]/cpx(n,0);
}
void solve(cpx a[],cpx b[],int n) //n項式相乘,FFT solving, ansult to a
{  //請確保n是2的次冪
    FFT(a,n,1); //DFT
    FFT(b,n,1);
    for(int i=0;i<n;i++)
        a[i]=a[i]*b[i]; //點值多項式相乘
    FFT(a,n,-1);//IDFT
}
 
cpx a[MAX];
int xp[MAX];
ll ans[MAX],arr[MAX];
int main()
{
    int n,p;
    scanf("%d%d",&n,&p);
    int g=getRoot(p);
    for(int i=0,x=1;i<p-1;i++)
    {
        xp[x]=i;
        x=(ll)x*g%p;
    }
    int len=p+p-2;
    for(;(len&-len)!=len;len+=len&-len);
    ll t=0;
    for(int i=1;i<=n;i++)
    {
        scanf("%lld",&arr[i]);
        if(arr[i]%p==0)t++;
        else a[xp[arr[i]%p]]+=cpx(1,0);
    }
    FFT(a,len,1);
    for(int i=0;i<len;i++)a[i]=a[i]*a[i];
    FFT(a,len,-1);
    for(int i=0,x=1;i<len;i++)
    {
        ans[x]+=ll(a[i].real()+0.5);
        x=(ll)x*g%p;
    }
    for(int i=1;i<=n;i++)
    {
        if(arr[i]>=p)puts("0");
        else if(arr[i]==0)printf("%lld\n",t*t+2*t*(n-t));
        else printf("%lld\n",ans[arr[i]]);
    }
}