2016 acm香港网络赛 A题. A+B Problem (FFT)

原题地址：https://open.kattis.com/problems/aplusb

FFT代码参考kuangbin的博客：http://www.cnblogs.com/kuangbin/archive/2013/07/24/3210565.html

A+B Problem

Given N integers in the range [−50000,50000], how many ways are there to pick three

integers ai, aj, ak, such that i, j, k are pairwise distinct and ai+aj=ak? Two ways

are different if their ordered triples (i,j,k)of indices are different.

Input

The first line of input consists of a single integer N

(1≤N≤200000). The next line consists of N space-separated integers a1,a2,…,aN

Output

Output an integer representing the number of ways.

Sample Input 1	Sample Output 1
4 1 2 3 4	4

Sample Input 2	Sample Output 2
6 1 1 3 3 4 6	10

Author(s): Tung Kam Chuen

Source: * Regional Online Preliminary 2016

题意：给一个数列，从中选三个数 ai, aj, ak，使得ai+aj=ak，问共有多少组（ i, j, k）满足条件。

其实就是FFT。

注意一下a数组的数据范围，a[i]可能为负数，所有a[i]+50000，把负数转化为正数处理;

如果数组里有0，先把0删掉，最后特殊处理。

把a数组转化成num数组，num[i]表示i出现的次数。

然后num数组和num数组卷积，得到新的num数组。

num数组意义：num[x]表示使ai+aj=x,（i,j）的取法有多少种。

#include <algorithm>
#include <cstring>
#include <string.h>
#include <iostream>
#include <list>
#include <map>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
#include <cstdio>
#include <cmath>

#define LL long long
#define N 200005
#define INF 0x3ffffff

using namespace std;

const double PI = acos(-1.0);

struct Complex // 复数
{
    double r,i;
    Complex(double _r = ,double _i = )
    {
        r = _r; i = _i;
    }
    Complex operator +(const Complex &b)
    {
        return Complex(r+b.r,i+b.i);
    }
    Complex operator -(const Complex &b)
    {
        return Complex(r-b.r,i-b.i);
    }
    Complex operator *(const Complex &b)
    {
        return Complex(r*b.r-i*b.i,r*b.i+i*b.r);
    }
};

void change(Complex y[],int len) // 二进制平摊反转置换 O(logn)
{
    int i,j,k;
    for(i = , j = len/;i < len-;i++)
    {
        if(i < j)swap(y[i],y[j]);
        k = len/;
        while( j >= k)
        {
            j -= k;
            k /= ;
        }
        if(j < k)j += k;
    }
}
void fft(Complex y[],int len,int on) //DFT和FFT
{
    change(y,len);
    for(int h = ;h <= len;h <<= )
    {
        Complex wn(cos(-on**PI/h),sin(-on**PI/h));
        for(int j = ;j < len;j += h)
        {
            Complex w(,);
            for(int k = j;k < j+h/;k++)
            {
                Complex u = y[k];
                Complex t = w*y[k+h/];
                y[k] = u+t;
                y[k+h/] = u-t;
                w = w*wn;
            }
        }
    }
    if(on == -)
        for(int i = ;i < len;i++)
            y[i].r /= len;
}

const int M =;          // a数组所有元素+M，使a[i]>=0
const int MAXN = ;

Complex x1[MAXN];
int a[MAXN/];                //原数组
long long num[MAXN];     //利用FFT得到的数组
long long tt[MAXN];        //统计数组每个元素出现个数

int main()
{
    int n=;                   // n表示除了0之外数组元素个数
    int tot;
    scanf("%d",&tot);
    memset(num,,sizeof(num));
    memset(tt,,sizeof(tt));

    int cnt0=;             //cnt0 统计0的个数
    int aa;

    for(int i = ;i < tot;i++)
        {
            scanf("%d",&aa);
            if(aa==) {cnt0++;continue;}  //先把0全删掉，最后特殊考虑0
            else a[n]=aa;
            num[a[n]+M]++;
            tt[a[n]+M]++;
            n++;
        }

    sort(a,a+n);
    int len1 = a[n-]+M+;
    int len = ;

    while( len < *len1 ) len <<= ;

    for(int i = ;i < len1;i++){
         x1[i] = Complex(num[i],);
    }
    for(int i = len1;i < len;i++){
         x1[i] =Complex(,);
    }
    fft(x1,len,);

    for(int i = ;i < len;i++){
        x1[i] = x1[i]*x1[i];
    }
    fft(x1,len,-);

    for(int i = ;i < len;i++){
         num[i] = (long long)(x1[i].r+0.5);
    }

    len = *(a[n-]+M);

    for(int i = ;i < n;i++) //删掉ai+ai的情况
           num[a[i]+a[i]+*M]--;
/*
   for(int i = 0;i < len;i++){
            if(num[i]) cout<<i-2*M<<' '<<num[i]<<endl;
    }
*/
    long long ret= ;

    int l=a[n-]+M;

    for(int i = ;i <=l; i++)   //ai,aj,ak都不为0的情况
        {
            if(tt[i]) ret+=(long long)(num[i+M]*tt[i]);
        }

    ret+=(long long)(num[*M]*cnt0);        // ai+aj=0的情况

    if(cnt0!=)
        {
            if(cnt0>=) {                   //ai,aj,ak都为0的情况
                long long tmp=;
                tmp*=(long long)(cnt0);
                tmp*=(long long)(cnt0-);
                tmp*=(long long)(cnt0-);
                ret+=tmp;
            }
             for(int i = ;i <=l; i++)
                {
                    if(tt[i]>=){             // x+0=x的情况
                        long long tmp=(long long)cnt0;
                        tmp*=(long long)(tt[i]);
                        tmp*=(long long)(tt[i]-);
                        ret+=tmp*;
                    }
                }
        }

    printf("%lld\n",ret);

    return ;
}