gym102920 H. Needle（FFT）

题意：

给定长度为n的序列a，长度为m的序列b，长度为p的序列c，
问有多少个a(i)+c(i)=2b(i)

数据范围：n,m,p<=5e4，-3e4<=a(i),b(i),c(i)<=3e4

解法：

对a和c列两个多项式,x^k的系数就是k出现的次数,
用fft相乘,设相乘之后x^k的系数为sum[k].

对于每个b[i],ans+=sum[b[i]*2]即可.

code：

#include <bits/stdc++.h>
using namespace std;
const int maxm=2e6+5;
const double P=acos(-1.0);
struct CC{//复数
    double x,y;
    CC(double xx=0,double yy=0){x=xx,y=yy;}
    CC operator+(const CC &a)const{return CC(x+a.x,y+a.y);}
    CC operator-(const CC &a)const{return CC(x-a.x,y-a.y);}
    CC operator*(const CC &a)const{return CC(x*a.x-y*a.y,x*a.y+y*a.x);}
};
int a[maxm];
int b[maxm];
int c[maxm];
int n,m,p;
CC x1[maxm],x2[maxm];
int sum[maxm];
int rev[maxm];
void fft(CC y[],int len,int on){//on为1或者-1,-1的时候表示逆变换
    for(int i=0;i<len;i++)if(i<rev[i])swap(y[i],y[rev[i]]);
    for(int h=2;h<=len;h<<=1){
        CC wn(cos(-on*2*P/h),sin(-on*2*P/h));
        for(int j=0;j<len;j+=h){
            CC w(1,0);
            for(int k=j;k<j+h/2;k++){
                CC u=y[k];
                CC 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;
        }
    }
}
signed main(){
    int maa=0,mab=0,mac=0;
    cin>>n;
    for(int i=1;i<=n;i++){
        cin>>a[i];
        a[i]+=3e4;
        maa=max(maa,a[i]);
    }
    cin>>m;
    for(int i=1;i<=m;i++){
        cin>>b[i];
        b[i]+=3e4;
        mab=max(mab,b[i]);
    }
    cin>>p;
    for(int i=1;i<=p;i++){
        cin>>c[i];
        c[i]+=3e4;
        mac=max(mac,c[i]);
    }
    int len=1,l=0;
    while(len<=maa+mac)len<<=1,l++;
    for(int i=0;i<len;i++){
        rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
    }
    for(int i=1;i<=n;i++){
        x1[a[i]].x+=1;
    }
    for(int i=1;i<=p;i++){
        x2[c[i]].x+=1;
    }
    fft(x1,len,1);
    fft(x2,len,1);
    for(int i=0;i<len;i++){
        x1[i]=x1[i]*x2[i];
    }
    fft(x1,len,-1);
    for(int i=0;i<=len;i++){
        sum[i]=(int)(x1[i].x+0.5);
    }
    len=n+m;
    int ans=0;
    for(int i=1;i<=m;i++){
        ans+=sum[b[i]*2];
    }
    cout<<ans<<endl;
    return 0;
}

---