A*B 原根+FFT优化

Posted 2020-10-05 HITLJR

题目地址 https://loj.ac/problem/6156

#include <bits/stdc++.h>
const long long mod = 1e9+7;
const double ex = 1e-10;
#define inf 0x3f3f3f3f
using namespace std;
const double PI = acos(-1.0);
const int MAXN = 60044*4;
struct Complex{
    double x,y;
    Complex(double _x = 0.0,double _y =0.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);
    }
};
void change (Complex y[],int len){
    int i,j,k;
    for (i = 1,j = len/2;i<len -1;i++){
        if (i<j) swap(y[i],y[j]);
        k = len/2;
        while (j >= k){
            j-=k;
            k/=2;
        }
        if (j<k) j+=k;
    }
}
void fft(Complex y[],int len, int on)
{
    change(y,len);
    for (int h = 2; h<=len;h<<=1){
        Complex wn(cos(-on*2*PI/h),sin(-on*2*PI/h));
        for (int j =0 ; j< len; j+=h){
            Complex w(1,0);
            for (int k = j;k < j+h/2;k++){
                Complex u = y[k];
                Complex t = w*y[k+h/2];
                y[k] = u+t;
                y[k+h/2] = u-t;
                w = w*wn;
            }
        }
    }
    if (on == -1)
        for (int i = 0; i<len; i++)
        y[i].x /= len;

}

Complex x1[MAXN];
long long x2[MAXN];
long long num[MAXN];
long long ans[MAXN/4];
long long quick_s(long long x,long long y,long long N){
    long long ans = 1;
    while (y){
        if (y % 2) ans = (ans * x) % N;
        x = (x*x) % N;
        y/=2;
    }
    return ans % N;
}
vector <long long> p;
int b[62345];
void getprime(int N){
    for (int i = 2;i<=N;i++){
        if ( b[i] == 0 ){
            p.push_back((long long)i);
            int j = 2;
            while (i*j <= N)
                b[i*j] = 1,j++;
        }
    }
}

int getroot(long long N){
    vector <long long> C;
    C.clear();
    long long t = N-1;
    for (long long i = 0; i<p.size() && p[i]*p[i] <= N; i++){
        if (t % p[i] == 0 && t!=1){
            C.push_back((N-1)/p[i]);
            while (t % p[i] == 0 && t!=1) t/=p[i];
        }
    }
    if (t!=1) C.push_back((N-1)/t);
    for (long long i = 2 ; i<=N; i++){
        int flag = 1;
        for (int j = 0 ; j< C.size() && flag ; j++ )
            if (quick_s(i,C[j],N) == 1&&C[j]!=N-1) flag = 0;
        if (flag && (quick_s(i,N-1,N) == 1)) return i;
    }
    return 0;
}
int main()
{
    getprime(60000);
    int T;
    scanf("%d",&T);
    while (T--){
        long long n,m;
        scanf("%lld%lld",&n,&m);
        long long  root = getroot(m);
        memset(num,0,sizeof(num));
        memset(x2,0,sizeof(x2));
        memset(ans,0,sizeof(ans));
        int len1 = m;
        for (int i = 1; i<=n; i++){
            int d;
            scanf("%d",&d);
            num[d%m]++;
        }
        int len = 1;
        while (len < len1*2) len <<=1;
        for (int i = 1; i<m;i++)
            x1[i] = Complex(num[quick_s(root,i,m)],0);
        for (int i = m ; i<len ; i++)
            x1[i] = Complex(0,0);
        x1[0] = Complex(0,0);
        fft(x1,len,1);
        for (int i = 0 ; i<len ; i++)
            x1[i] = x1[i] * x1[i];
        fft(x1,len,-1);

        ans[0] = num[0] * (num[0]-1)/2 + num[0] * (n-num[0]);
        printf("%d\n",ans[0]);
        for (int i = 0; i<=2*m;i++)
            x2[i] = (long long)(x1[i].x + 0.5);
        for (long long i = 1; i<m; i++){
            long long t = quick_s(root,i,m);
            x2[2*i] -= num[t] * num[t];
            x2[2*i-1]/=2;
            x2[2*i]/=2;
            x2[2*i] += (num[t] * (num[t] - 1)/2);
        }
        if (m==2){
            printf("%lld\n",n*(n-1)/2-ans[0]);
        }
        else{
            for (long long  i = 1; i<m ; i++){
                ans[quick_s(root,i,m)] = x2[i]+ x2[m-1+i];
            }
            for (int i = 1; i<m ; i++){
                printf("%lld\n",ans[i]);
            }
        }
    }
    return 0;
}

View Code