Kattis A+B Problem(FFT,计数)

来源:互联网 发布:淘宝女装代销货源 编辑:程序博客网 时间:2024/04/30 07:47

题目链接

题意:给出n个数,(i,j,k)i,j,k互不相同,使得a[i]+a[j]=a[k],问这样的三元组有几个。(HK网赛的第一题,当时xjb写的一直wa,今天学FFT,突然遇见他)
其实就是计算数组的a+b的值的个数,这个就是卷积了,FFT搞一下。
题目的数据是有负数的,加上一个偏移,全部变为非负的,一开始的时候把0单独分开,最后在统计0的贡献。

#include<cstdio>#include<cstring>#include<algorithm>#include<cmath>using namespace std;#define cl(a,b) memset(a,b,sizeof(a))#define ll long long#define pb push_back#define gcd __gcdconst double EPS = 1e-8;const int maxn = 8e5+1000;const int inf  = 0x3f3f3f3f;const double PI = acos(-1.0);struct Complex{    double r,i;    Complex(double _r=0.0,double _i=0.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);    }};/*进行FFT和IFFT前的反转变换位置i 和 i的二进制反转表示的数 位置交换len 必须是2的幂*/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]);//交换下标互为反转的元素,i<j保证只交换一次        k=len/2;//i做正常的加1,j做反转的加一,始终保持i,j是反转的        while(j>=k){            j-=k;k/=2;        }        if(j<k)j+=k;    }}/*    做FFT len必须是2的幂    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].r/=len;}int offset=50000;int a[maxn/4];ll num[maxn],tmp[maxn];Complex x[maxn];int main(){    int N;    int n=0;ll cnt0=0;scanf("%d",&N);    int mx=-offset-100;    for(int i=0;i<N;i++){        int x;scanf("%d",&x);        if(x){            x+=offset;            a[n++]=x;            num[x]++;tmp[x]++;            mx=max(mx,x);        }        else cnt0++;    }    int len=1;    while(len<mx*2)len<<=1;    for(int i=0;i<len;i++)x[i]=Complex(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=0;i<len;i++)num[i]=(ll)(x[i].r+0.5);    ll ans=0;    for(int i=0;i<n;i++){        num[a[i]+a[i]]--;    }    for(int i=0;i<=mx;i++){        ans+=tmp[i]*num[i+offset];    }    ans+=num[offset<<1]*cnt0;//a+b=0    if(cnt0>=3)ans+=cnt0*(cnt0-1)*(cnt0-2);//0+0=0    for(int i=0;i<=mx;i++){        if(tmp[i]>=2){            ans+=2*cnt0*(tmp[i]-1)*tmp[i];//x+0=x        }    }    printf("%lld\n",ans);    return 0;}
0 0