16年香港网络赛 A+B Problem FFT

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>

using namespace std;
//FFT模板
const double PI = acos(-1.0);
struct Complex{
	double x,y;
	Complex(double _x = 0.0,double _y = 0.0){
		x = _x;
		y = _y;
	}
	Complex operator-(const Complex &b)const{
		return Complex(x - b.x,y - b.y);
	}
	Complex operator+(const Complex &b)const{
		return Complex(x + b.x,y + b.y);
	}
	Complex operator*(const Complex &b)const{
		return Complex(x*b.x - y*b.y,x*b.y + y*b.x);
	}
};
/*
*进行FFT和IFFT前的反置变换
*位置i和i的二进制反转后的位置互换
*len必须为2的幂
*/
void change(Complex y[],int len){
	int i,j,k;
	for(int i = 1,j = len/2;i<len-1;i++){
		if(i < j)	swap(y[i],y[j]);
		//交换互为小标反转的元素，i<j保证交换一次
		//i做正常的+1，j做反转类型的+1，始终保持i和j是反转的
		k = len/2;
		while(j >= k){
			j = j - k;
			k = k/2;
		}
		if(j < k)	j+=k;
	}
}
/*
*做FFT
*len必须是2^k形式
*on == 1时是DFT，on == -1时是IDFT
*/
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];
				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].x /= len;
		}
	}
}

const int MAXN = 200000;
Complex x1[2*MAXN+5],x2[2*MAXN+5];
long long sum[MAXN+5],a[MAXN+5];
long long ans[2*MAXN+5];
int T = 5e4;
int n,cnt0;
int main(){
	scanf("%d",&n);
	cnt0=0;
	for(int i=0;i<n;i++){
        scanf("%lld",&a[i]);
        if(a[i] == 0) cnt0++;
        sum[a[i]+T]++;
	}
    //cout << cnt0 <<endl;
	int len=1;
    while(len < MAXN) len <<= 1;
    for(int i=0;i<MAXN;i++){
        x1[i] = x2[i] = Complex(1.0*sum[i],0.0);
    }
   // for(int i=n;i<len;i++){
     //   x1[i] = Complex(0,0);
   // }
//    cout << "ok" <<endl;
//    FFT(x1,len,1);
//    cout << "ok" <<endl;
//    for(int i =0;i<n;i++){
//        x2[i] = x1[i]*x1[i];
//    }
//    FFT(x2,len,-1);
    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++){
        ans[i] = (long long)(x1[i].x+0.5);
    }
    for(int i=0;i<n;i++){
        ans[(a[i]+T)*2]--;
    }
//    for(int i=0;i<=10;i++)
//        cout << ans[a[i]+2*T] << "  ";

    long long cnt = 0;
    for(int i=0;i<n;i++){
        cnt += ans[a[i]+2*T];
       // cout << cnt <<endl;
        cnt -= (cnt0-(a[i]==0))*2;
    }
    printf("%lld\n",cnt);


	return 0;
}