Kattis A B Problem(FFT)

Kattis A B Problem(FFT)
目錄


題意:求n個數中任取三個組合成ai aj=ak 的對數
思路:把給的數作為多項式的次冪,出現的次數為係數,然後用FFT進行多項式的乘法

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <algorithm>
#define maxn (1<<20)
#define pi 3.141592653589793238462643383
#define M 50000
using namespace std;
int p[1000100],n,zero=0,ma=0;
long long ans=0;
struct complex
{
double re,im;
complex(double r=0.0,double i=0.0) {re=r,im=i;}
void print() {printf("%.lf ",re);}
} a[maxn*2],b[maxn*2],W[2][maxn*2],c[maxn*2];
int N,na,nb,rev[maxn*2];
complex operator  (const complex&A,const complex&B) {return complex(A.re B.re,A.im B.im);}
complex operator -(const complex&A,const complex&B) {return complex(A.re-B.re,A.im-B.im);}
complex operator *(const complex&A,const complex&B) {return complex(A.re*B.re-A.im*B.im,A.re*B.im A.im*B.re);}
void FFT(complex*a,int f)
{
complex x,y;
for (int i=0; i<N; i  ) if (i<rev[i]) swap(a[i],a[rev[i]]);
for (int i=1; i<N; i<<=1)
for (int j=0,t=N/(i<<1); j<N; j =i<<1)
for (int k=0,l=0; k<i; k  ,l =t) x=W[f][l]*a[j k i],y=a[j k],a[j k]=y x,a[j k i]=y-x;
if (f) for (int i=0; i<N; i  ) a[i].re/=N;
}
void work()
{
for (int i=0; i<N; i  )
{
int x=i,y=0;
for (int k=1; k<N; x>>=1,k<<=1) (y<<=1)|=x&1;
rev[i]=y;
}
for (int i=0; i<N; i  ) W[0][i]=W[1][i]=complex(cos(2*pi*i/N),sin(2*pi*i/N)),W[1][i].im=-W[0][i].im;
}
void init()
{   memset(a,0,sizeof(a));
memset(b,0,sizeof(b));
scanf("%d",&n); 
for (int i=0; i<n; i  ) 
{
scanf("%d",&p[i]);
if(!p[i]) zero  ;
a[p[i] M].re  ;
ma=max(p[i] M,ma); 
}
for(int i=0;i<=ma;i  ) b[i].re=a[i].re;
for (N=1; N<ma; N<<=1); N<<=1;
}
void doit()
{
work(),FFT(a,0),FFT(b,0);
for (int i=0; i<N; i  ) a[i]=a[i]*b[i];
FFT(a,1);
for(int i=0;i<n;i  )  a[2*(p[i] M)].re--;
double ans=0;
for(int i=0;i<n;i  )  ans =a[p[i] 2*M].re;
ans-=2*(long long)zero*(long long )(n-1);
if(ans<0) printf("0");
else printf("%.lf\n",ans);
}
int main()
{   
freopen("1.txt","r",stdin);
freopen("1.out","w",stdout); 
init();
doit();
}