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: Hong Kong 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 = 0,double _i = 0)
{
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 = 1, j = len/2;i < len-1;i++)
{
if(i < j)swap(y[i],y[j]);
k = len/2;
while( j >= k)
{
j -= k;
k /= 2;
}
if(j < k)j += k;
}
}
void fft(Complex y[],int len,int on) //DFT和FFT
{
change(y,len);
for(int h = 2;h <= len;h <<= 1)
{
Complex wn(cos(-on*2*PI/h),sin(-on*2*PI/h));
for(int j = 0;j < len;j += h)
{
Complex w(1,0);
for(int k = j;k < j+h/2;k++)
{
Complex u = y[k];
Complex t = w*y[k+h/2];
y[k] = u+t;
y[k+h/2] = u-t;
w = w*wn;
}
}
}
if(on == -1)
for(int i = 0;i < len;i++)
y[i].r /= len;
}

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

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

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

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

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

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

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

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

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

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

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

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

int l=a[n-1]+M;

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

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

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

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

return 0;
}