BZOJ3771 : Triple

FFT后减去重复选取的方案数。

#include<cstdio>
#include<cmath>
#include<algorithm>
using namespace std;
typedef long long ll;
const int N=131072;
int n,i,m,k;ll ans[N],sum[N],f1[N],f2[N];
struct comp{
double r,i;comp(double _r=0,double _i=0){r=_r;i=_i;}
comp operator+(const comp x){return comp(r+x.r,i+x.i);}
comp operator-(const comp x){return comp(r-x.r,i-x.i);}
comp operator*(const comp x){return comp(r*x.r-i*x.i,r*x.i+i*x.r);}
}a[N],b[N],c[N];
const double pi=acos(-1.0),eps=1e-4;
void FFT(comp a[],int n,int t){
  for(int i=1,j=0;i<n-1;i++){
    for(int s=n;j^=s>>=1,~j&s;);
    if(i<j)swap(a[i],a[j]);
  }
  for(int d=0;(1<<d)<n;d++){
    int m=1<<d,m2=m<<1;
    double o=pi/m*t;comp _w(cos(o),sin(o));
    for(int i=0;i<n;i+=m2){
      comp w(1,0);
      for(int j=0;j<m;j++){
        comp &A=a[i+j+m],&B=a[i+j],t=w*A;
        A=B-t;B=B+t;w=w*_w;
      }
    }
  }
  if(t==-1)for(int i=0;i<n;i++)a[i].r/=n;
}
inline void read(int&a){char c;while(!(((c=getchar())>='0')&&(c<='9')));a=c-'0';while(((c=getchar())>='0')&&(c<='9'))(a*=10)+=c-'0';}
int main(){
  read(n);
  while(n--){
    read(i),a[i].r+=1,f1[i]++,ans[i]++;
    if(i>m)m=i;
  }
  for(m=m*3+1,k=1;k<m;k<<=1);
  FFT(a,k,1);
  for(i=0;i<k;i++)b[i]=a[i]*a[i];
  FFT(b,k,-1);
  for(i=0;i<k;i++)sum[i]=(ll)(b[i].r+0.5);
  for(i=0;i<k;i+=2)sum[i]-=f1[i>>1];
  for(i=0;i<k;i++)ans[i]+=sum[i]>>1;
  for(i=0;i<k;i++)c[i]=a[i]*a[i]*a[i];
  FFT(c,k,-1);
  for(i=0;i<k;i++)sum[i]=(ll)(c[i].r+0.5);
  for(i=0;i<k;i++)b[i]=comp(0,0);
  for(i=0;i<k/2;i++)b[i<<1]=comp(f1[i],0);
  FFT(b,k,1);
  for(i=0;i<k;i++)c[i]=a[i]*b[i];
  FFT(c,k,-1);
  for(i=0;i<k;i++)f2[i]=(ll)(c[i].r+0.5);
  for(i=0;i<k;i+=3)f2[i]-=f1[i/3],sum[i]-=f1[i/3];
  for(i=0;i<k;i++)sum[i]-=f2[i]*3;
  for(i=0;i<k;i++)ans[i]+=sum[i]/6;
  for(i=0;i<k;i++)if(ans[i])printf("%d %lld
",i,ans[i]);
  return 0;
}