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();
}