[LOJ #6156][找原根][FFT || NTT] A*B Problem

题目要求对于所有0≤k<m的k都输出解，可以考虑ansk的生成函数。

ansk=∑i×j=kai∗aj

ai为i的出现次数，这样可以矩乘优化，但是复杂度太大，因为m是质数，所以依然可以考虑用原根，把ai转化成gbi，那么

ansk=∑i+j=kbi∗bj

就可以用FFT完成了。

#include <cstdio>#include <iostream>#include <algorithm>#include <cmath>using namespace std;const int N=600010;const double PI=acos(-1);int n,m,G,cnt,num,mi[N],a[N],b[N],rev[N];long long ans[N];struct E{  double real,imag;  E(double a=0,double b=0):real(a),imag(b){}  friend E operator+(E a,E b){ return E(a.real+b.real,a.imag+b.imag); }  friend E operator-(E a,E b){ return E(a.real-b.real,a.imag-b.imag); }  friend E operator*(E a,E b){    return E(a.real*b.real-a.imag*b.imag,a.real*b.imag+b.real*a.imag);  }}c[N*6],w[2][N*6];inline int Pow(int x,int y){  int ret=1;  for(;y;y>>=1,x=1LL*x*x%m) ret=1LL*x*ret%m;  return ret;}bool isprim(int g){  for(int i=1;i<m;i++) mi[i]=0;  for(int i=1,ig=g;i<=m-2;i++,ig=1LL*ig*g%m) mi[ig]=i;  for(int i=2;i<m;i++) if(!mi[i]) return false;  return true;}inline void Getg(){  for(G=2;G<=m;G++)    if(isprim(G)) break;}inline void Pre(int n){  num=n;  E g(cos(2*PI/num),sin(2*PI/num));  w[0][0]=w[1][0]=1;  for(int i=1;i<num;i++) w[1][i]=w[1][i-1]*g;  for(int i=1;i<num;i++) w[0][i]=w[1][num-i];}inline void FFT(E *a,int n,int r){  for(int i=0;i<n;i++) if(rev[i]>i) swap(a[i],a[rev[i]]);  for(int i=1;i<n;i<<=1)    for(int j=0;j<n;j+=(i<<1))      for(int k=0;k<i;k++){    E x=a[j+k],y=w[r][num/(i<<1)*k]*a[j+k+i];    a[j+k]=x+y; a[j+k+i]=x-y;      }  if(!r) for(int i=0;i<n;i++) a[i].real/=n;}int main(){  int t; scanf("%d",&t);  while(t--){    scanf("%d%d",&n,&m); Getg(); cnt=0;    if(m==2){      int cnt1=0,cnt0=0;      for(int i=1;i<=n;i++){    scanf("%d",&a[i]);    if(a[i]&1) cnt1++;    else cnt0++;      }      printf("%lld\n%lld\n",1LL*cnt0*cnt1+1LL*cnt0*(cnt0-1)/2,1LL*cnt1*(cnt1-1)/2);      continue;    }    for(int i=0;i<m;i++) b[i]=0;    for(int i=1;i<=n;i++){      scanf("%d",&a[i]);      if(a[i]%m==0) cnt++;      else b[mi[a[i]%m]]++;    }    for(int i=0;i<m;i++) c[i]=b[i];    printf("%lld\n",1LL*cnt*(n-cnt)+1LL*cnt*(cnt-1)/2);    int M,L=0; for(M=1;M<=(m<<1);M<<=1,L++); Pre(M<<1);    for(int i=m;i<M<<1;i++) c[i]=0;    for(int i=1;i<M<<1;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<L);    FFT(c,M<<1,1);    for(int i=0;i<M<<1;i++) c[i]=c[i]*c[i];    FFT(c,M<<1,0);    for(int i=1;i<m;i++) ans[i]=(long long)(c[mi[i]].real+c[mi[i]+m-1].real+0.1);    for(int i=1;i<=n;i++)      if(a[i]%m) ans[1LL*a[i]*a[i]%m]--;    for(int i=1;i<m;i++) printf("%lld\n",ans[i]/2);  }  return 0;  }