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

4
1 2 3 4

Sample Output 1

4

Sample Input 2

6
1 1 3 3 4 6

Sample Output 2

10

思路

这道题出的真的很好$!!!!$明白$fft$的很好的题目
这里要结合$fft$的性质-解决$ai+aj=ak$的问题,因为$fft$中任意两项乘积,$x的k1和x的k2$指数次方相加,所以可以构建 系数为$ai$的次数,幂为$ai$的大小的多项式,两个多项式相乘,对于每一个$ak$,它的系数即为次数,不过中间有一点小细节需要处理~那就是出现0相加的时候会导致结果和另一个加数相等,不和要求,所以要减掉$ai==0||aj==0$的情况(为0的时候少减一个)
$(Ps:$推荐一个良性博客,讲的很详细$fft解决A+B$$)$

代码如下

#include<cstdio>
#include<cmath>
#include<algorithm>
#include<cstring>
#include<iostream>
#define N 200005
typedef long long ll;
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 &rhs){
        return Complex(r+rhs.r,i+rhs.i);
    }
    Complex operator - (const Complex &rhs){
        return Complex(r-rhs.r,i-rhs.i);
    }
    Complex operator * (const Complex &rhs){
        return Complex(r*rhs.r-i*rhs.i,i*rhs.r+r*rhs.i);
    }
}F[N*4];
ll a[N*4],num[N*4];
ll zero,cnt;
int len;

 
inline void sincos(double theta,double &p0,double &p1){//找根
    p0=sin(theta);
    p1=cos(theta);
}

void FFT(Complex P[], int n, int oper){ //fft的基本操作
    for(int i=1,j=0;i<n-1;i++){
        for(int s=n;j^=s>>=1,~j&s;);
        if(i<j)swap(P[i],P[j]);
    }
    Complex unit_p0;
    for(int d=0;(1<<d)<n;d++){
        int m=1<<d,m2=m*2;
        double p0=pi/m*oper;
        sincos(p0,unit_p0.i,unit_p0.r);
        for(int i=0;i<n;i+=m2){
            Complex unit=1;
            for(int j=0;j<m;j++){
                Complex &P1=P[i+j+m],&P2=P[i+j];
                Complex t=unit*P1;
                P1=P2-t;
                P2=P2+t;
                unit=unit*unit_p0;
            }
        }
    }
    if(oper==-1)for(int i=0;i<len;i++)P[i].r/=len;
}

void Conv(Complex a[],int len){//求卷积
    FFT(a,len,1);//FFT
    for(int i=0;i<len;++i)F[i]=F[i]*F[i];
    FFT(F,len,-1);//IFFT
}

int main(){
	    len=1;zero=0;cnt=0;
     	memset(num,0,sizeof(num));
        int n;
        scanf("%d",&n);
        for(int i=0;i<n;i++)
        {
	     scanf("%lld",&a[i]);
	     if(a[i]==0) //统计0的数量
	     zero++;
	     num[a[i]+50000]++;//统计系数
        }
        sort(a,a+n);
        int len1=a[n-1]+1+50000;
        while(len<len1*2)//得到长度,必须为2^n
        len<<=1;
        
        for(int i=0;i<len1;i++)//赋值
        F[i]=Complex(num[i],0);
        
        for(int i=len1;i<len;i++)
        F[i]=Complex(0,0);
        
        Conv(F,len);//fft
        
        for(int i=0;i<len;i++)
        {
		  num[i]=(ll)(F[i].r+0.5);//得到新的系数
	    }
        
        for(int i=0;i<n;i++)
        num[a[i]+a[i]+2*50000]--;//首先减掉自身和自身的
        
        for(int i=0;i<n;i++)
        {
        	cnt+=num[a[i]+2*50000];
        	if(a[i]!=0)//不为0的时候减掉2倍
			cnt-=zero*2;
		    else
		    cnt-=(zero-1)*2;
		}
		printf("%lld\n",cnt); 
		return 0;
}