bzoj3509 [CodeChef] COUNTARI 分块+FFT

Description

---

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

第一行一个整数N（N<=10^5）。
接下来一行N个数A[i]（A[i]<=30000）。

Solution

---

我会n^2logn做法！枚举中点和左端点主席树右端点！于是可以用这个来拍

拆一下柿子可以发现要求的是满足$2a_j=a_i+a_k$的数量
一个比较naive的做法就是我们枚举j，然后把左边的桶和右边的桶卷积的第2*a[j]位就是j对答案的贡献
考虑对序列分块。我们把块外左边的桶和块外右边的桶FFT，枚举块内作为中点
注意到这样会算漏一个端点和中点在同一块、另一个端点在块外的情况，我们用一个桶暴力算即可。这里的暴力需要一点奇迹银壳不然会被卡shi

一开始8s+，然后各种卡常。包括但不限于：调整块的大小、循环展开、手写复数类、回滚清空数组
然后我就发现，手写复数类之后反而更慢了quq

Code

---

#include <stdio.h>
#include <string.h>
#include <math.h>
#include <algorithm>
#include <complex>
#include <map>
#define rep(i,st,ed) for (register int i=st,_=ed;i<=_;++i)
#define drp(i,st,ed) for (register int i=st,_=ed;i>=_;--i)
#define fill(x,t) memset(x,t,sizeof(x))

typedef std:: complex <double> com;
typedef long long LL;
const double pi=3.141592653589;
const int N=200005;
const int M=1048578;

int rev[N],bel[N],st[N],ed[N],B;
int a[N],cnt[N],map[N];

com b[M],c[M],d[M];

int read() {
    int x=0,v=1; char ch=getchar();
    for (;ch<'0'||ch>'9';v=(ch=='-')?(-1):(v),ch=getchar());
    for (;ch<='9'&&ch>='0';x=x*10+ch-'0',ch=getchar());
    return x*v;
}

void FFT(com *a,int len,double f) {
    for (int i=1;i<len;i++) if (i<rev[i]) std:: swap(a[i],a[rev[i]]);
    for (int i=1;i<len;i<<=1) {
        com wn(cos(pi/i),f*sin(pi/i));
        for (int j=0;j<len;j+=i*2) {
            com w(1,0);
            for (int k=0;k<i;k++) {
                com u=a[j+k],v=a[j+k+i]*w;
                a[j+k]=u+v; a[j+k+i]=u-v;
                w*=wn;
            }
        }
    }
    if (f==-1) for (int i=1;i<len;i++) a[i]/=len;
}

int main(void) {
    freopen("data.in","r",stdin);
    freopen("myp.out","w",stdout);
    int n=read(),max=0; B=sqrt(n*300);
    rep(i,1,n) {
        a[i]=read(); cnt[a[i]]++;
        max=std:: max(max,a[i]);
        bel[i]=(i-1)/B+1;
        if (!st[bel[i]]) st[bel[i]]=i;
        ed[bel[i]]=i;
    }
    int len,lg; for (len=1,lg=0;len<=max*2;len<<=1,lg++) ;
    rep(i,1,len) rev[i]=(rev[i/2]/2)|((i&1)<<(lg-1));
    LL ans=0;
    rep(i,1,bel[n]) {
        rep(j,st[i],ed[i]) {
            rep(k,j+1,ed[i]) {
                int tmp=a[j]+a[j]-a[k];
                ans+=map[tmp];
            }
            ++map[a[j]];
        }
    } rep(i,0,max) map[i]=0;
    drp(i,bel[n],1) {
        drp(j,ed[i],st[i]) {
            drp(k,j-1,st[i]) {
                int tmp=a[j]+a[j]-a[k];
                ans+=map[tmp];
            }
            map[a[j]]+=1;
        }
    } rep(i,0,max) map[i]=0;
    rep(i,1,bel[n]) {
        rep(j,st[i],ed[i]) {
            rep(k,j+1,ed[i]) {
                int tmp=a[j]+a[j]-a[k];
                ans-=map[tmp];
            }
            ++map[a[j]];
        }
        rep(j,st[i],ed[i]) --map[a[j]];
    }
    rep(i,1,bel[n]) {
        rep(j,0,len) b[j]=c[j]=0;
        rep(j,1,st[i]-1) b[a[j]]+=1;
        rep(j,ed[i]+1,n) c[a[j]]+=1;
        FFT(b,len,1); FFT(c,len,1);
        rep(i,0,len) d[i]=b[i]*c[i];
        FFT(d,len,-1);
        rep(j,st[i],ed[i]) {
            ans+=(d[a[j]+a[j]].real()+0.5);
        }
    }
    printf("%lld\n", ans);
    return 0;
}