1. 程式人生 > >牛客國慶集訓派對Day1 I Steins;Gate(原根 + FFT)

牛客國慶集訓派對Day1 I Steins;Gate(原根 + FFT)

上一次用到原根這個東西,應該是一年之前了吧……

所謂原根,就是指對於某個數字P,滿足它的原根g,g的0~P-2次冪在模P的意義下互不相同,或者說g的1~P-1次冪在模P的意義下無不相同。一般來說,原根只能夠進行列舉求解,但是由於原根一般較小(2或者3),所以可以接受。

對於本題,要求對於每一個數字ak計算在模P意義下,有多少個ai*aj等於ak。直接做顯然是O(N^2)的,不能夠滿足條件。看到這個題目形式和資料範圍,很容易往FFT等方向上想,但是這裡的ai的範圍是1e9級別,而且這個是乘法,直接FFT也是不行的。由於P滿足是質數,所以我們考慮用原根。

我們令g為質數P的原根,那麼對於一個數字ai,唯一存在一個數字bi,使得 g^{b_i}=a_i

。那麼我們把所有的ai用這種形式表示,於是對於原本的式子ai*aj≡ak(mod P),有g^{b_i}*g^{b_j}=g^{b_i+b_j}=g^{b_k},那麼可以有bi+bj=bk。這樣問題就從乘法變成了加法,FFT就可以排上用場了。而且你會發現,這裡的bi的大小都是在P以內的。這樣直接FFT即可,最後特殊處理以下0的情況即可。具體見程式碼:

#include<bits/stdc++.h>
#define PI 3.1415926535
#define eps 1e-8
#define mod 998244353
#define LL long long
#define pb push_back
#define lb lower_bound
#define ub upper_bound
#define INF 0x3f3f3f3f
#define sf(x) scanf("%d",&x)
#define sc(x,y,z) scanf("%d%d%d",&x,&y,&z)
#define clr(x,n) memset(x,0,sizeof(x[0])*(n+5))
#define file(x) freopen(#x".in","r",stdin),freopen(#x".out","w",stdout)
using namespace std;

const int N = 500010;

struct Complex
{
    double r,i;
    Complex(double real=0.0,double image=0.0)
    {
        r=real; i=image;
    }

    Complex operator +(const Complex o){return Complex(r+o.r,i+o.i);}
    Complex operator -(const Complex o){return Complex(r-o.r,i-o.i);}
    Complex operator *(const Complex o){return Complex(r*o.r-i*o.i,r*o.i+i*o.r);}
} b[N];

namespace FFT
{
    int len;

    void brc(Complex *y, int l)
    {
        register int i,j,k;
        for( i = 1, j = l / 2; i < l - 1; i++)
        {
            if (i < j) swap(y[i], y[j]);
            k = l / 2; while ( j >= k) j -= k,k /= 2;
            if (j < k) j += k;
        }
    }

    void FFT(Complex *y, int len, double on)
    {
        register int h, j, k;
        Complex u, t; brc(y, len);
        for(h = 2; h <= len; h <<= 1)
        {
            Complex wn(cos(on * 2 * PI / h), sin(on * 2 * PI / h));
            for(j = 0; j < len; j += h)
            {
                Complex w(1, 0);
                for(k = j; k < j + h / 2; k++)
                {
                    u = y[k]; t = w * y[k + h / 2];
                    y[k] = u + t; y[k + h / 2] = u - t;
                    w = w * wn;
                }
            }
        }
        if (on<0) for (int i = 0; i < len; i++) y[i].r/=len;
    }

    void multiply(Complex *A,int lenA)
    {
        for(len = 1; len < 2*lenA - 1; len <<= 1);
        for (int i = lenA; i < len; i++) A[i] = 0;
        FFT(A,len ,1 );
        for (int i = 0;i < len; i++) A[i] = A[i] * A[i];
        FFT(A, len, -1);
    }
}

int a[N],id[N],n,P,g;
LL ans[N];

inline bool check(int g)
{
    LL tmp=1;
    for(int i=1;i<P-1;i++)
    {
        tmp=tmp*g%P;
        if (tmp==1) return 0;
    }
    return 1;
}

int main()
{
    sf(n); sf(P);
    for(int i=2;;i++)
        if (check(i)) {g=i;break;}
    LL tmp=1,t=0; id[1]=0;
    for(int i=1;i<P-1;i++)
    {
        tmp=tmp*g%P;
        id[tmp]=i;
    }

    for(int i=1;i<=n;i++)
    {
        sf(a[i]);
        if (a[i]%P==0) t++;
        else b[id[a[i]%P]].r++;
    }
    FFT::multiply(b,P-1);
    for(int i=0;i<P*2;i++)
        ans[i%(P-1)]+=(LL)(b[i].r+0.5);
    for(int i=1;i<=n;i++)
    {
        if (a[i]>=P) puts("0");
        else if (a[i]==0) printf("%lld\n",t*(n-t)*2+t*t);
        else printf("%lld\n",ans[id[a[i]]]);
    }

    return 0;
}