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

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

$$ans_k=\sum_{i\times j=k} a_i*a_j$$
$a_i$为 $i$的出现次数，这样可以矩乘优化，但是复杂度太大，因为m是质数，所以依然可以考虑用原根，把$a_i$转化成 $g^{b_i}$，那么

$$ans_k=\sum_{i+j=k} b_i*b_j$$ 就可以用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;  
}