BZOJ4836: [Lydsy1704月赛]二元运算-分治FFT

传送门

题意：

定义二元运算 opt 满足

$$x \ opt \ y =\left\{\begin{aligned}x+y&,&x<y \\x-y&,&x≥y\end{aligned}\right.$$

现在给定一个长为 n 的数列 a 和一个长为 m 的数列 b ，接下来有 q 次询问。每次询问给定一个数字 c

你需要求出有多少对 (i, j) 使得 $a_i\ opt\ b_j=c$ 。

Solution：

首先我们抛开限制来看：

如果只考虑加法的话，对于每个询问k，$ans=\sum_{i=0}^ka_i*b_{k-i}$

只考虑减法的话，对于每个询问k，$ans=\sum_{i=k}^na_i*b_{i-k}$

我们将a数组翻转后可以得到$ans=\sum_{i=k}^na_{n-i}*b_{i-k}=\sum_{i=0}^{n-k}a_i*b_{n-k-i}$

但是，由于有大小限制，所以说不能简单地FFT

对于加法的情况，可能会计算到$x≥y$的情况，而对于减法则不会

所以说我们只需要对加法情况进行分治FFT，对减法情况只需要一次FFT即可求出所有答案

关于分治FFT，具体就是每次对于$[l,r]$这段区间，处理$a[l,mid]$和$b[mid+1,r]$即可

代码：

#include<cstdio>
#include<cstring>
#include<iostream>
#include<cmath>
using namespace std;
const double pi=acos(-1);
int n,m,T,q,a[50010],b[50010],maxn;
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);
    }
}a1[2*65540],a2[2*65540];
long long ans[200010];
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]);
        k=len/2;
        while (j>=k) j-=k,k>>=1;
        if (j<k) j+=k;
    }
}
void FFT(complex y[],int len,int ifi)
{
    change(y,len);
    for (int h=2;h<=len;h<<=1)
    {
        complex wn(cos(2*pi*ifi/h),sin(2*pi*ifi/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 v=w*y[k+h/2];
                y[k]=u+v;y[k+h/2]=u-v;
                w=w*wn;
            }
        }
    }
    if (ifi==-1) for (int i=0;i<len;i++) y[i].x/=len;
}
void solve(int l,int r)
{
    int mid=l+r>>1;
    if (l==r) {//ans[0]+=a[l]*b[r];
    return;}
    solve(l,mid);solve(mid+1,r);
    int len=1;
    while (len<=r-l+1) len<<=1;
    for (int i=l;i<=mid;i++) a1[i-l]=complex(a[i],0);
    for (int i=mid-l+1;i<len;i++) a1[i]=complex(0,0);
    for (int i=mid+1;i<=r;i++) a2[i-mid-1]=complex(b[i],0);
    for (int i=r-mid;i<len;i++) a2[i]=complex(0,0);
    FFT(a1,len,1);FFT(a2,len,1);
    for (int i=0;i<len;i++) a1[i]=a1[i]*a2[i];
    FFT(a1,len,-1);
    for (int i=0;i<len;i++) ans[i+l+mid+1]+=(long long)(a1[i].x+0.5);
}
int main()
{
//  freopen("4836.in","r",stdin);
//    freopen("4836.out","w",stdout);
    scanf("%d",&T);
    while (T--)
    {
        memset(a,0,sizeof(a));
        memset(b,0,sizeof(b));
        memset(ans,0,sizeof(ans));
        scanf("%d%d%d",&n,&m,&q);
        for (int x,i=1;i<=n;i++) scanf("%d",&x),a[x]++,maxn=max(maxn,x);
        for (int x,i=1;i<=m;i++) scanf("%d",&x),b[x]++,maxn=max(maxn,x);
        solve(0,maxn);
        int len=1;
        while (len<=2*maxn) len<<=1;
        for (int i=0;i<=maxn;i++) a1[maxn-i]=complex(a[i],0);
        for (int i=maxn+1;i<len;i++) a1[i]=complex(0,0);
        for (int i=0;i<=maxn;i++) a2[i]=complex(b[i],0);
        for (int i=maxn+1;i<len;i++) a2[i]=complex(0,0);
    //  for (int i=0;i<len;i++) printf("%d ",(long long)(a1[i].x+0.5));cout<<endl;
    //  for (int i=0;i<len;i++) printf("%d ",(long long)(a2[i].x+0.5));cout<<endl;
        FFT(a1,len,1);FFT(a2,len,1);
        for (int i=0;i<len;i++) a1[i]=a1[i]*a2[i];
        FFT(a1,len,-1);
        //for (int i=0;i<len;i++) printf("%d ",(long long)(a1[i].x+0.5));cout<<endl;
        for (int i=0;i<=maxn;i++) ans[maxn-i]+=(long long)(a1[i].x+0.5);
        for (int x,i=1;i<=q;i++)
            scanf("%d",&x),printf("%lld\n",ans[x]);
    }
}