SPOJ-Triple Sums 3个不同位置的数字之和 fft

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/liufengwei1/article/details/86629787

WA了一晚上，不知道为撒

然后今天下午换了种方法，又调了一个小时，最后发现fft的len搞错了。。。

算出3次项的fft可以直接x1[i]=x1[i]*x1[i]*x1[i];这中间包括i=j=k的情况，和i=j || j=k || i=k的情况是非法的，直接ans3[3*(a[i]+M)]吧第一种非法情况去掉

第二种非法情况用x2[2*(a[i]+M)]和x1进行fft，那么就能得出2*ai+aj的所有情况ans2[]，又因为第二种情况在有3种排列，所以合法的情况为ans3[i]-ans2[i]*3.

最后的合法情况中又存在i,j,k3数的全排列，所以最后要除以6

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
#define maxl 240010

using namespace std;

const int M=20000;
const double pi=acos(-1.0);
int n,mx,len;
int a[maxl];
long long ans1[maxl],ans2[maxl],ans3[maxl];

struct complex
{
	double r,i;
	complex(double r1=0.0,double i1=0.0)
	{
		r=r1;i=i1;
	}
	complex operator + (const complex &b)
	{
		return complex(r+b.r,i+b.i);
	}
	complex operator - (const complex &b)
	{
		return complex(r-b.r,i-b.i);
	}
	complex operator * (const complex &b)
	{
		return complex(r*b.r-i*b.i,r*b.i+i*b.r);
	}
}x1[maxl],x2[maxl],x3[maxl];

inline void prework()
{
	scanf("%d",&n);
	mx=0;
	memset(x1,0,sizeof(x1));
	for(int i=1;i<=n;i++)
	{
		scanf("%d",&a[i]);
		if(a[i]+M>mx)
			mx=a[i]+M;
	}
}
void change(complex y[],int len)
{
	int i,j;
	for(i=1,j=len/2;i<len-1;i++)
	{
		if(i<j)
			swap(y[i],y[j]);
		int k=len/2;
		while(j>=k)
			j-=k,k/=2;
		if(j<k)
			j+=k;
	}
}

void fft(complex y[],int len,int on)
{
	change(y,len);
	for(int l=2;l<=len;l<<=1)
	{
		complex wn(cos(-on*2*pi/l),sin(-on*2*pi/l));
		for(int j=0;j<len;j+=l)
		{
			complex w(1,0);
			for(int k=j;k<j+l/2;k++)
			{
				complex u=y[k];
				complex t=w*y[k+l/2];
				y[k]=u+t;
				y[k+l/2]=u-t;
				w=w*wn;
			}
		}
	}
	if(on==-1)
		for(int i=0;i<len;i++)
			y[i].r/=len;
}

inline void mainwork()
{
	len=1;
	while(len<mx*3+1)
		len<<=1;
	for(int i=0;i<len;i++)
		x1[i]=complex(0.0,0.0);
	for(int i=1;i<=n;i++)
		x1[a[i]+M].r+=1.0;
	fft(x1,len,1);
	for(int i=0;i<len;i++)
		x1[i]=x1[i]*x1[i]*x1[i];
	fft(x1,len,-1);
	for(int i=0;i<len;i++)
		ans3[i]=(long long)(x1[i].r+0.5);
	for(int i=1;i<=n;i++)
		ans3[3*(a[i]+M)]--;

	for(int i=0;i<len;i++)
		x1[i]=x2[i]=complex(0.0,0.0);
	for(int i=1;i<=n;i++)
		x1[a[i]+M].r+=1.0,x2[2*(a[i]+M)].r+=1.0;
	fft(x1,len,1);
	fft(x2,len,1);
	for(int i=0;i<len;i++)
		x2[i]=x1[i]*x2[i];
	fft(x2,len,-1);
	for(int i=0;i<len;i++)
		ans2[i]=(long long)(x2[i].r+0.5);
	for(int i=1;i<=n;i++)
		ans2[3*(a[i]+M)]--;

	
	for(int i=0;i<len;i++)
		ans3[i]=(ans3[i]-ans2[i]*3)/6;
	
}

inline void print()
{
	for(int i=0;i<len;i++)
	if(ans3[i]>0)
		printf("%d : %lld\n",i-3*M,ans3[i]);
}

int main()
{
	prework();
	mainwork();
	print();
	return 0;
}