[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) 的个数。

思路

http://blog.csdn.net/qpswwww/article/details/44228381

代码

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

void rev(Complex a[],int n)
{
    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>>=1;
        }
        if(j<k) j+=k;
    }
}

void FFT(Complex a[],int n,int flag)
{
    rev(a,n);
    for(int i=1;i<n;i<<=1)
    {
        Complex wn=Complex(cos(PI/i),flag*sin(PI/i));
        for(int j=0;j<n;j+=(i<<1))
        {
            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];
                a[j+k]=x+y,a[j+k+i]=x-y;
            }
        }
    }
    if(flag==-1)
        for(int i=0;i<n;i++)
            a[i]=Complex(a[i].real()/n,a[i].imag());
}

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

int main()
{
    int n;
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        int x;
        scanf("%d",&x);
        x+=20000;
        a[x].real()+=1;
        b[x+x].real()+=1;
        C[x+x+x]++;
    }
    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[i])/6;
        if(ans>0)
            printf("%d : %lld\n",i-60000,ans);
    }
    return 0;
}