BZOJ 3509 分块FFT

//By SiriusRen
#include <cmath>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const double pi=acos(-1);
const int N=100050;
int n,nn,num[N],R[N],L,Block,block[N],cnt[50][N];
long long ans;
struct Complex{
    double x,y;Complex(){}
    Complex(double X,double Y){x=X,y=Y;}
}A[N],B[N],C[N];
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);}
Complex operator/(Complex a,int b){return Complex(a.x/b,a.y/b);}
void FFT(Complex *a,int f){
    for(int i=0;i<n;i++)if(i<R[i])swap(a[i],a[R[i]]);
    for(int i=1;i<n;i<<=1){
        Complex wn=Complex(cos(pi/i),f*sin(pi/i));
        for(int j=0;j<n;j+=(i<<1)){
            Complex w=Complex(1,0);
            for(int k=0;k<i;k++,w=w*wn){
                Complex x=a[j+k],y=w*a[j+k+i];
                a[j+k]=x+y,a[j+k+i]=x-y;
            }
        }
    }
    if(!~f)for(int i=0;i<n;i++)a[i]=a[i]/n;
}
int main(){
    scanf("%d",&nn);
    for(int i=1;i<=nn;i++)scanf("%d",&num[i]);
    for(n=1;n<=60000;n<<=1)L++;
    for(int i=0;i<n;i++)R[i]=(R[i>>1]>>1)|((i&1)<<(L-1));
    Block=min(int(sqrt(nn)*10),nn);
    for(int i=1;i<=nn;i++)block[i]=(i-1)/Block+1;
    for(int i=1;i<=nn;i++)cnt[block[i]][num[i]]++;
    for(int I=1;I<=block[nn];I++){
        int L=lower_bound(block+1,block+1+nn,I)-block,R=upper_bound(block+1,block+1+nn,I)-block-1;
        for(int j=L;j<=R;j++){
            cnt[I][num[j]]--;
            for(int i=L;i<j;i++)
                if(num[j]*2-num[i]>=0)ans+=cnt[I][num[j]*2-num[i]];
        }
    }
    for(int i=1;i<=nn;i++)cnt[0][num[i]]++;
    for(int I=1;I<=block[nn];I++){
        int L=lower_bound(block+1,block+1+nn,I)-block,R=upper_bound(block+1,block+1+nn,I)-block-1;
        for(int i=L;i<=R;i++)cnt[0][num[i]]--;
        for(int j=L;j<=R;j++)
            for(int i=L;i<j;i++)
                if(num[j]*2-num[i]>=0)ans+=cnt[0][num[j]*2-num[i]];
    }
    for(int i=1;i<=nn;i++)cnt[0][num[i]]++;
    for(int I=block[nn];I;I--){
        int L=lower_bound(block+1,block+1+nn,I)-block,R=upper_bound(block+1,block+1+nn,I)-block-1;
        for(int i=L;i<=R;i++)cnt[0][num[i]]--;
        for(int k=L;k<=R;k++)
            for(int j=k-1;j>=L;j--)
                if(num[j]*2-num[k]>=0)ans+=cnt[0][num[j]*2-num[k]];
    }
    for(int I=1;I<=block[nn];I++){
        for(int i=0;i<n;i++)A[i].x=A[i].y=B[i].x=B[i].y=0;
        int L=lower_bound(block+1,block+1+nn,I)-block,R=upper_bound(block+1,block+1+nn,I)-block-1;
        for(int i=1;i<L;i++)A[num[i]].x++;
        for(int i=R+1;i<=nn;i++)B[num[i]].x++;
        FFT(A,1),FFT(B,1);
        for(int i=0;i<n;i++)C[i]=A[i]*B[i];
        FFT(C,-1);
        for(int i=L;i<=R;i++)ans+=(long long)(C[num[i]*2].x+0.2);
    }
    printf("%lld\n",ans);
}