SPOJ8372 TSUM - Triple Sums 生成函数 FFT 容斥

题目链接

题意：

给你n个整数，求从其中任选3个所有可能组成的数是哪些和能组成的这些数分别有多少种可能的组合，我们认为相同的3个数的不同排列是同一种，每个数只能用一次。

题解：

假如不考虑每个数只能用一次这个限制，并且也先不管不同排列算作同一种，这个不同排列只要最后除一个排列数就好了。那么我们只需要构造一个生成函数 A ( x ) A(x) A(x)， A ( x ) = a 1 x + a 2 x 2 + . . . + a n x n A(x)=a_1x+a_2x^2+...+a_nx^n A(x)=a1​x+a2​x2+...+an​xn，其中 a n a^n an表示这 n n n个数中有几个数的值是 n n n，答案是 A ∗ A ∗ A A*A*A A∗A∗A。

下面我们考虑同一个数会被多次用到的方案数，用一点容斥的思想来算这个数。我们先减去三个中有两个数是用了同一个数的方案数，我们设为 B B B函数， B = a 1 x 2 + a 2 x 4 + . . . + a n x 2 n B=a_1x^2+a_2x^4+...+a_nx^{2n} B=a1​x2+a2​x4+...+an​x2n。我们用 A ∗ B A*B A∗B即可求出有两个数相同的方案数，其中包含了只有两个数相同和三个数全相同的情况。由于我们先不考虑不同的排列是同一种的情况，于是我们这里先乘上一个 C 3 2 C_{3}^2 C32​，但是我们发现其实我们把三个数全相同的多减了两次，于是要加上两倍的三个全相同的情况。我们构造 C C C函数为三个数全相同的方案数， C = a 1 x 3 + a 2 x 6 + . . . + a n x 3 n C=a_1x^3+a_2x^6+...+a_nx^{3n} C=a1​x3+a2​x6+...+an​x3n。最后我们还要除以一个3个数的全排列6。于是答案就是 A ∗ A ∗ A − 3 ∗ A ∗ B + 2 ∗ C 6 \frac{A*A*A-3*A*B+2*C}{6} 6A∗A∗A−3∗A∗B+2∗C​。其中的多项式乘法用FFT优化。

我们发现有负数，那么我们把负数集体加上20000即可。由于我们最后乘了3次方，所以答案的下标要减60000。

代码：

#include <bits/stdc++.h>
using namespace std;

int n,m,l,rev[3000010];
complex<double> a[3000010],b[3000010],c[3000010];
const double pi=acos(-1);
inline void fft(complex<double> *a,int dft)
{
	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)
	{
		complex<double> wn(cos(pi/i),dft*sin(pi/i));
		for(int p=i<<1,j=0;j<n;j+=p)
		{
			complex<double> w(1,0);
			for(int k=0;k<i;++k)
			{
				complex<double> x=a[j+k],y=w*a[i+j+k];
				a[j+k]=x+y;
				a[i+j+k]=x-y;
				w=w*wn;
			}
		}
	}
}
int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;++i)
	{
		int x;
		scanf("%d",&x);
		a[x+20000]+=(1);
		b[2*(x+20000)]+=(1);
		c[3*(x+20000)]+=(1);
		m=max(m,3*(x+20000));
	}
	m++;
	m<<=1;
	for(n=1;n<=m;n<<=1)
	++l;
	for(int i=0;i<n;++i)
	rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
	fft(a,1);
	fft(b,1);
	fft(c,1);
	complex<double> ji1=(6),ji2=(3),ji3=(2);
	for(int i=0;i<n;++i)
	a[i]=(a[i]*a[i]*a[i]-ji2*a[i]*b[i]+ji3*c[i])/ji1;
	fft(a,-1);
	for(int i=0;i<m;++i)
	{
		long long ji=(long long)(a[i].real()/n+0.5);
		if(ji>0)
		printf("%d : %lld\n",i-60000,ji);
	}
	return 0;
}