Asia Hong Kong Online Preliminary A (FFT)

题目链接：点击这里

题意：给出n个数字（∈[−50000,50000]）,求有多少(i,j,k)满足ai+aj=ak(要求i,j,k互不相同)。

0比较麻烦直接拎出来记录个数，其他的加一个50000避免负数直接存起来。然后开一个数组记录每一个数字出现多少次，这个数组和自己做一次卷积即可。答案分成几部分：

1. i+j=k: 直接遍历扔掉0的数组，把算完卷积的结果扔进去；

2. 0+0=0：直接用0的个数可以得出结果；

3. x+y=0：直接用卷积中0次项的系数 或者暴力遍历；

4. 0+x=x/x+0=x：暴力遍历统计。

#include <cstdio>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <queue>
#include <map>
using namespace std;
#define pi acos (-1)
#define maxn 511111
#define add 50000

struct plex {
double x, y;
plex (double _x = 0.0, double _y = 0.0) : x (_x), y (_y) {}
plex operator + (const plex &a) const {
return plex (x+a.x, y+a.y);
}
plex operator - (const plex &a) const {
return plex (x-a.x, y-a.y);
}
plex operator * (const plex &a) const {
return plex (x*a.x-y*a.y, x*a.y+y*a.x);
}
}x[maxn];

void change (plex *y, int len) {
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(plex y[],int len,int on)
{
change(y,len);
for(int h = 2; h <= len; h <<= 1)
{
plex wn(cos(-on*2*pi/h),sin(-on*2*pi/h));
for(int j = 0;j < len;j+=h)
{
plex w(1,0);
for(int k = j;k < j+h/2;k++)
{
plex u = y[k];
plex 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].x /= len;
}

int n, cnt, zero;
int a[maxn], num[maxn];

long long solve () {
long long ans = 0;
int len = 1;
while (len < 200000) {
len <<= 1;
}
for (int i = 0; i < len; i++) {
x[i] = plex (num[i], 0);
}
fft (x, len, 1);
for (int i = 0; i < len; i++) {
x[i] = x[i]*x[i];
}
fft (x, len, -1);

for (int i = 1; i <= cnt; i++) {//i+j=k
long long tmp = (long long) (x[a[i]+add*2].x + 0.5);
ans += tmp;
}
for (int i = 1; i <= cnt; i++) {//减去i+i=j
if (a[i]%2 == 0)
ans -= num[a[i]/2+add];
}

if (zero >= 3) ans += 1LL*zero*(zero-1)*(zero-2);//0+0=0
for (int i = 0; i <= add*2; i++) {//x+0=y y+0=x
if (num[i] > 1) {
ans += 1LL*2*num[i]*(num[i]-1)*zero;
}
}
ans += 1LL*(long long)(x[2*add].x+0.5)*zero;//x+y=0

return ans;
}

int main () {
scanf ("%d", &n);
zero = cnt = 0;
memset (num, 0, sizeof num);
for (int i = 1; i <= n; i++) {
scanf ("%d", &a[++cnt]);
if (a[cnt] == 0) {
zero++;
cnt--;
continue;
}
num[a[cnt]+add]++;
}
long long ans = solve ();
printf ("%lld\n", ans);
return 0;
}
/*
10
-50000 -50000 -50000 0 0 0 50000 50000 50000 1
4
0 0 -1 1
*/