BZOJ 4836: [Lydsy1704月赛]二元运算 分治FFT

Code:

#include<bits/stdc++.h>
#define ll long long
#define maxn 500000
#define setIO(s) freopen(s".in","r",stdin)
using namespace std;
const double pi=acos(-1.0);
struct cpx
{
    double x,y;
    cpx(double a=0,double b=0) {x=a,y=b; }
    cpx operator+(const cpx b) { return cpx(x+b.x,y+b.y); }
    cpx operator-(const cpx b) { return cpx(x-b.x,y-b.y); }
    cpx operator*(const cpx b) { return cpx(x*b.x-y*b.y,x*b.y+y*b.x); }
}A[maxn],B[maxn];
void FFT(cpx *a,int n,int flag)
{
    for(int i=0,k=0;i<n;++i)
    {
        if(i>k) swap(a[i],a[k]);
        for(int j=n>>1;(k^=j)<j;j>>=1);
    }
    for(int mid=1;mid<n;mid<<=1)
    {
        cpx wn(cos(pi/mid), flag*sin(pi/mid)),x,y;
        for(int i=0;i<n;i+=(mid<<1))
        {
            cpx w(1,0);
            for(int j=0;j<mid;++j,w=w*wn) x=a[i+j],y=w*a[i+j+mid],a[i+j]=x+y,a[i+j+mid]=x-y;
        }
    }
    if(flag==-1) for(int i=0;i<n;++i) a[i].x/=(double)n;
}
int n,m;
int arr[maxn],brr[maxn],bucka[maxn],buckb[maxn];
ll answer[maxn];
void solve(int l,int r)
{
    if(l==r)  { answer[0]+=1ll*bucka[l]*buckb[l]; return; }
    int mid=(l+r)>>1,len,a=0,b=0;
    for(len=1;len<=(r-l+1);len<<=1);
    for(int i=0;i<=len;++i) A[i].x=B[i].x=A[i].y=B[i].y=0;
    for(int i=l;i<=mid;++i) A[a++].x=bucka[i];
    for(int i=mid+1;i<=r;++i) B[b++].x=buckb[i];
    FFT(A,len,1),FFT(B,len,1);
    for(int i=0;i<len;++i) A[i]=A[i]*B[i];
    FFT(A,len,-1);
    for(int i=mid+1+l;i<=mid+r;++i) answer[i]+=(ll)(A[i-mid-1-l].x+0.5);
    for(int i=0;i<=len;++i) A[i].x=A[i].y=B[i].x=B[i].y=0;
    for(int i=mid+1;i<=r;++i) A[i-mid-1].x=bucka[i];
    for(int i=l;i<=mid;++i) B[mid-i]=buckb[i];
    FFT(A,len,1),FFT(B,len,1);
    for(int i=0;i<len;++i) A[i]=A[i]*B[i];
    FFT(A,len,-1);
    for(int i=0;i<=r-l;++i) answer[i+1]+=(ll)(A[i].x+0.5);
    solve(l,mid), solve(mid+1,r);
}
void work()
{
    memset(bucka,0,sizeof(bucka));
    memset(buckb,0,sizeof(buckb));
    memset(answer,0,sizeof(answer));
    int q,Max=0;
    scanf("%d%d%d",&n,&m,&q);
    for(int i=1;i<=n;++i) scanf("%d",&arr[i]),Max=max(Max,arr[i]);
    for(int i=1;i<=m;++i) scanf("%d",&brr[i]),Max=max(Max,brr[i]);
    for(int i=1;i<=n;++i) bucka[arr[i]]++;
    for(int i=1;i<=m;++i) buckb[brr[i]]++;
    solve(0,Max);
    for(int i=1;i<=q;++i)
    {
        int x;
        scanf("%d",&x);
        printf("%lld\n",answer[x]);
    }
}
int main()
{
    // setIO("input");
    int T;
    scanf("%d",&T);
    for(int i=1;i<=T;++i) work();
    return 0;
}