Kattis aplusb A+B Problem(FFT)

A+B Problem
Given NN integers in the range [−50000,50000], how many ways are there to pick three integers aiai, ajaj, akak, such that ii, jj, kk are pairwise distinct and ai+aj=akai+aj=ak? Two ways are different if their ordered triples (i,j,k)(i,j,k) of indices are different.

Input
The first line of input consists of a single integer NN (1≤N≤200000). The next line consists of NN space-separated integers a1,a2,…,aNa1,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

题意

求n个数中任取三个组合成$a_i+a_j=a_k$对数

思路

这题和HDU4609一样,做法也一样,但是因为数据范围里会出现负数所以我们要给每个数加上一个M=50000
一样的做法,把给的数作为多项式的次幂,次数为系数,然后用FFT进行多项式的乘法

注意1
这题(i,j)和(j,i)认为不同的,所以答案不需要除二，但是自己取自己的情况不包括在内所以还要历遍数组减去两倍的每一个数
注意2
可能存在给的数为0的情况,当你给的数为0是就会有$0+0=0$的情况和$a_i+0=a_i$的情况,因为自己取自己是非法的,所以这里要删去一些情况
我们先统计一下0的个数为z个，那么就有
$2z(n-z)+2z(z-1)=2z(n-1)$
所以最后的答案再减去$2z(n-1)$就好了
注意3
因为每个数都加上了M
所以那个题意中的式子应为$a_i+M+a_j+M=a_k+M+M$

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <math.h>
#include <stack>
#define PI acos(-1)
using namespace std;
typedef long long ll;
const int MAXN = 1e6+5;
struct complex {
    double a, b;
    complex(double aa = 0.0, double bb = 0.0) { a = aa; b = bb; }
    complex operator +(const complex &e) { return complex(a + e.a, b + e.b); }
    complex operator -(const complex &e) { return complex(a - e.a, b - e.b); }
    complex operator *(const complex &e) { return complex(a * e.a - b * e.b, a * e.b + b * e.a); }
}x1[MAXN], x2[MAXN], x[MAXN];
void change(complex *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>>=1; }
        if(j < k) j += k;
    }
}
void fft(complex *y, int len, int on){
    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], 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].a /= len;
}
long long a[MAXN];
long long ans[MAXN];
long long M=50000;
int main()
{
    long long n;
    scanf("%lld",&n);
        memset(x1,0,sizeof(x1));
        memset(ans,0,sizeof(ans));
        long long maxn=0;
        long long zero=0;
        for(int i=0;i<n;i++)
        {
            scanf("%lld",&a[i]);
            if(!a[i]) zero++;
            x1[a[i]+M].a++;
            maxn=max(a[i]+M,maxn);
        }
        int len=1;
        while(len<2*maxn) len<<=1;
        fft(x1,len,1);
        memset(x,0,sizeof(x));
        for(int i=0;i<len;i++)
            x[i]=x1[i]*x1[i];
        fft(x,len,-1);
        for(int i=0;i<len;i++)
            ans[i]=(ll)(x[i].a+0.5);
        for(int i=0;i<n;i++)
            ans[2*(a[i]+M)]--;
        ll sum=0;
        for(int i=0;i<n;i++)
            sum+=ans[a[i]+M+M];
        sum-=2*zero*(n-1);
        printf("%lld\n",sum);
    return 0;
}