分块FFT--bzoj3509: [CodeChef] COUNTARI

传送门

给定一个长度为N的数组A[]，求有多少对i, j, k（1<=i<j<k<=N）满足A[k]-A[j]=A[j]-A[i]。

暴力的话就考虑枚举$j$，然后算两边的可以用$FFT$来算，但这样复杂度是$n\ast mx\ast logmx$，优化的话就可以用分块，而且块不能太小，然后块外用$FFT$，块内就暴力。

注意卡常
被卡了无数次以后发现不可以预处理要边做边处理

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#define maxn 30005
#define maxm 100005
using namespace std;
const double Pi=acos(-1.0);
 
inline int rd(){
    int x=0,f=1;char c=' ';
    while(c<'0' || c>'9') f=c=='-'?-1:1,c=getchar();
    while(c<='9' && c>='0') x=x*10+c-'0',c=getchar();
    return x*f;
}
inline int max(int x,int y){return x>y?x:y;}
inline int min(int x,int y){return x<y?x:y;}
 
struct complex{
    double x,y;
    complex(double xx=0,double yy=0) {x=xx,y=yy;}
}a[maxn<<2],b[maxn<<2];
complex operator +(complex a,complex b){return complex(a.x+b.x,a.y+b.y);}
complex operator -(complex a,complex b){return complex(a.x-b.x,a.y-b.y);}
complex operator *(complex a,complex b){return complex(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}

int n,cntl[maxn],cntr[maxn],limit=1,l,rev[maxn<<2];
int block,L[300],R[300],tot,mx,val[maxm];
long long res;
 
inline void FFT(complex *F,int type){
    for(int i=0;i<limit;i++) if(i<rev[i]) swap(F[i],F[rev[i]]);
    for(int mid=1;mid<limit;mid<<=1){
        complex Wn(cos(Pi/mid),1.0*type*sin(Pi/mid));
        for(int r=mid<<1,j=0;j<limit;j+=r){
            complex w(1,0);
            for(int k=0;k<mid;k++,w=w*Wn){
                complex x=F[j+k],y=w*F[j+mid+k];
                F[j+k]=x+y,F[j+mid+k]=x-y;
            }
        }
    }
    if(type==-1){
        for(int i=0;i<limit;i++) F[i].x=F[i].x/limit;
    }
}
 
inline void prework(){
    for(int i=1;i<=tot;i++) L[i]=R[i-1]+1,R[i]=i*block;
    if(R[tot]<n) ++tot,L[tot]=R[tot-1]+1,R[tot]=n;
    while(limit<=2*mx) limit<<=1,++l;
    for(int i=0;i<limit;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
}
 
inline void solve(){
    for(int i=1;i<=tot;i++){
        for(int j=0;j<limit;j++) a[j].x=a[j].y=b[j].x=b[j].y=0;
        for(int j=L[i];j<=R[i];j++) cntr[val[j]]--;
        for(int j=0;j<=mx;j++) a[j].x=cntl[j],b[j].x=cntr[j];
    	FFT(a,1); FFT(b,1);
        for(int j=0;j<limit;j++) a[j]=a[j]*b[j];
        FFT(a,-1);
        for(int j=L[i];j<=R[i];j++) {
            res+=(int)(a[val[j]<<1].x+0.5);
            for(int k=L[i];k<=R[i];k++){
                if(k<j){
                    int x=2*val[j]-val[k];
                    if(x>=0 && x<=mx) res+=cntr[x];
                }
                else if(k>j){
                    int x=2*val[j]-val[k];
                    if(x>=0 && x<=mx) res+=cntl[x];
                }
            }
            cntl[val[j]]++;
        }
    }
}

int main(){
    n=rd();
    for(int i=1;i<=n;i++) val[i]=rd(),mx=max(mx,val[i]),cntr[val[i]]++;
    block=min((int)6*sqrt(n),n); tot=n/block;
    prework(); solve();
    printf("%lld\n",res);
    return 0;
}