欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

Kattis A+B Problem(FFT)

程序员文章站 2022-05-22 19:23:48
...

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


相关标签: FFT