SPOJ TSUM Triple Sums（FFT + 容斥）

题目

Source

http://www.spoj.com/problems/TSUM/

Description

You're given a sequence s of N distinct integers.
Consider all the possible sums of three integers from the sequence at three different indicies.
For each obtainable sum output the number of different triples of indicies that generate it.

Constraints:
N <= 40000, |si| <= 20000

Input

The first line of input contains a single integer N.
Each of the next N lines contain an element of s.

Output

Print the solution for each possible sum in the following format:
sum_value : number_of_triples

Smaller sum values should be printed first.

Sample Input

5
-1
2
3
0
5

Sample Output

1 : 1
2 : 1
4 : 2
5 : 1
6 : 1
7 : 2
8 : 1
10 : 1

 

分析

题目大概说给n个数，从中选出三个数求和，问能到的各个和分别有几种取法能够得到？

 

这题很容易。。因为刚做过HDU4609。。
就是根据初始的序列构造出三个一样的多项式，指数表示数字，系数表示该数字出现次数。
然后三个多项式的乘积相当于表示有顺序有放回地取数字的结果。这个用FFT求。

 

不过这不是组合，可以用容斥原理去掉那些取法重复的。
即减去3种两个取同一边的情况，这个也用FFT求；然后加上三个都取同一边的情况，for一遍即可求；最后除以3的阶乘。

 

注意到时间好像比较紧，所以我做了些处理，比如三个多项式相乘直接三个点值相乘、避免重复的DFT过程。。
。。然后没想到居然暂时列第二，与第一同时间：

 

 

代码

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;
#define MAXN 131072
const double PI=acos(-1.0);
 
struct Complex{
	double real,imag;
	Complex(double _real=0,double _imag=0):real(_real),imag(_imag){}
	Complex operator+(const Complex &cp) const{
		return Complex(real+cp.real,imag+cp.imag);
	}
	Complex operator-(const Complex &cp) const{
		return Complex(real-cp.real,imag-cp.imag);
	}
	Complex operator*(const Complex &cp) const{
		return Complex(real*cp.real-imag*cp.imag,real*cp.imag+cp.real*imag);
	}
	void setValue(double _real=0,double _imag=0){
		real=_real; imag=_imag;
	}
};
 
int len;
Complex wn[MAXN+1],wn_anti[MAXN+1];
 
void FFT(Complex y[],int op){
	for(int i=1,j=len>>1,k; i<len-1; ++i){
		if(i<j) swap(y[i],y[j]);
		k=len>>1;
		while(j>=k){
			j-=k;
			k>>=1;
		}
		if(j<k) j+=k;
	}
	for(int h=2; h<=len; h<<=1){
		Complex Wn=(op==1?wn[h]:wn_anti[h]);
		for(int i=0; i<len; i+=h){
			Complex W(1,0);
			for(int j=i; j<i+(h>>1); ++j){
				Complex u=y[j],t=W*y[j+(h>>1)];
				y[j]=u+t;
				y[j+(h>>1)]=u-t;
				W=W*Wn;
			}
		}
	}
	if(op==-1){
		for(int i=0; i<len; ++i) y[i].real/=len;
	}
}

Complex A[MAXN],B[MAXN];
double ans[MAXN];
int s[40100],cnt[80100];

int main(){
	for(int i=0; i<=MAXN; ++i){
		wn[i].setValue(cos(2.0*PI/i),sin(2.0*PI/i));
		wn_anti[i].setValue(wn[i].real,-wn[i].imag);
	}
	
	int n;
	scanf("%d",&n);
	for(int i=0; i<n; ++i){
		scanf("%d",&s[i]);
		s[i]+=20000;
		++cnt[s[i]];
	}
	
	len=MAXN;
	
	for(int i=0; i<=40000; ++i){
		B[i].setValue(cnt[i]);
	}
	FFT(B,1);
	for(int i=0; i<MAXN; ++i){
		A[i]=B[i]*B[i]*B[i];
	}
	FFT(A,-1);
	for(int i=0; i<MAXN; ++i){
		ans[i]=A[i].real;
	}
	
	memset(cnt,0,sizeof(cnt));
	for(int i=0; i<n; ++i){
		++cnt[s[i]+s[i]];
	}
	for(int i=0; i<=80000; ++i){
		A[i].setValue(cnt[i]);
	}
	for(int i=80001; i<MAXN; ++i){
		A[i].setValue(0);
	}
	FFT(A,1);
	for(int i=0; i<MAXN; ++i){
		A[i]=A[i]*B[i];
	}
	FFT(A,-1);
	for(int i=0; i<MAXN; ++i){
		ans[i]-=3*A[i].real;
	}
	
	for(int i=0; i<n; ++i){
		++ans[s[i]+s[i]+s[i]];
	}
	
	for(int i=0; i<MAXN; ++i){
		long long tmp=(long long)(ans[i]/6.0+0.5);
		if(tmp){
			printf("%d : %lld\n",i-60000,tmp);
		}
	}
	return 0;
}