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