[SPOJ TSUM]Triple Sums(FFT+生成函数+容斥原理)

题目链接

http://acm.hust.edu.cn/vjudge/problem/viewProblem.action?id=23842

题目大意

给出 n 个数字 a1,a2,a3...an ,对于所有存在的 S ，求和为 S 的三个数 ai,aj,ak 的有序对 (i,j,k) (i<j<k) 的个数。

思路

如果这个题没有限制三个数必须不同的话，非常简单，用一个多项式 A(x) ，其中 xi 的系数是数字为 i 的数的个数， A(x)3 中每个 xSk 的系数就是和为 Sk 的三个可以相同的数的个数。
但是现在加了三个数互不相同的限制条件，考虑固定三个数字之和 S 的情况，此时和为 S 的三个互不相同的数的组成方案数=所有的随便三个数(无限制条件)构成和为 S 的有序对个数-两个数相同的构成和为 S 的有序对个数+三个数都相同的构成和为 S 的有序对个数，这是非常显然的容斥原理。

用生成函数来表示的话(以下的sigma部分省略了这些 xi 前面的系数)
所有的随便三个数(无限制条件)的有序对构成的答案多项式 =A(x)3=∑x+3∑xy−6∑xyz
两个数相同的有序对构成的答案的多项式 =A(x)T(x2)=∑x+∑x2y
三个数都相同的有序对构成的答案多项式 =G(x3)=∑x3
其中的 x、y、z 代表次数不同的项

T(x2) 就是对于 xi 而言，它前面的系数是数字为 i2 的数的个数。显然 A(x)T(x2) 代表的就是三个数中2个相同，另一个不同，这样的答案多项式， G(x3) 类似，就是代表了三个数全部一样的答案多项式

那么我们最后要求的答案多项式显然是 ∑xyz=A(x)3−A(x)T(x2)+G(x3)6 ，反正做完FFT的求值过后，多项式的乘法就非常轻松了，完成上面的式子完全无压力

下面的代码我换了个求比特反转的做法，感觉好写也好理解一些，以后就用下面的当模板了。

代码

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>
#include <complex>

#define MAXN 1<<17
#define PI 3.1415926535897384626

using namespace std;

typedef long long int LL;
typedef complex<long double> Complex;

struct Polynomial
{
    int len;
    int c[MAXN]; //c[i]=x^i项的系数
}A,B,C;

void reverse(Complex a[],int n) //对长度为的复数序列a进行DFT和IDFT之前必要的反转变换操作(比特反转),n一定是2的幂次
{
    for(int i=1,j=n/2,k;i<n-1;i++)
    {
        if(i<j) swap(a[i],a[j]);
        k=n/2;
        while(j>=k)
        {
            j-=k;
            k/=2;
        }
        if(j<k) j+=k;
    }
}

void FFT(Complex a[],int n,int flag) //flag=1是求值，flag=-1是插值
{
    reverse(a,n);
    for(int i=1;i<n;i<<=1) //i=该层的相邻结点之间编号之差，logn次迭代模拟，自底向上
    {
        Complex wn=Complex(cos(PI/i),flag*sin(PI/i));
        for(int j=0;j<n;j+=(i<<1)) //编号为j的结点
        {
            Complex w=Complex(1,0);
            for(int k=0;k<i;k++,w=w*wn)
            {
                Complex x=a[j+k],y=w*a[j+k+i]; //x和y是当前节点左右递归下去后得到的fft值，当然，y(右儿子，奇数项)的fft值要乘上一个w
                a[j+k]=x+y;
                a[j+k+i]=x-y;
            }
        }
    }
    if(flag==-1) //如果是插值，还得在前面加个1/n的系数
        for(int i=0;i<n;i++)
            a[i].real()/=n; //!!!!
}

Complex a[MAXN],b[MAXN],c[MAXN];

int main()
{
    A.len=B.len=C.len=MAXN;
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        int x;
        scanf("%d",&x);
        x+=20000;
        A.c[x]++;
        B.c[x+x]++;
        C.c[x+x+x]++;
    }
    for(int i=0;i<MAXN;i++)
        a[i]=A.c[i],b[i]=B.c[i];
    FFT(a,MAXN,1);
    FFT(b,MAXN,1);
    for(int i=0;i<MAXN;i++)
        c[i]=a[i]*(a[i]*a[i]-3.0l*b[i]);
    FFT(c,MAXN,-1);
    for(int i=0;i<MAXN;i++)
    {
        LL ans=((LL)(c[i].real()+0.5)+2*C.c[i])/6;
        if(ans>0) //和为i-60000的情况是存在的
            printf("%d : %lld\n",i-60000,ans);
    }
    return 0;
}