BZOJ[3509][CodeChef] COUNTARI 分块+FFT

传送门ber~

因为$A[k]-A[j]=A[j]-A[i]$，即$2A_j=A_k+A_i$
所以我们可以对每个位置的左右分别搞一个生成函数，卷积起来统计$2A_j$次项的系数就可以了
时间复杂度$\varTheta(nm\log m)$，过不去

可以考虑分块，块内的贡献暴力解决，块外的贡献生成函数，然后枚举块内每一个$2*A_i$次项的系数加在一起
时间复杂度$\varTheta(BlockNum*m\log m+BlockNum*M)$

代码如下：

#include<algorithm>
#include<ctype.h>
#include<cstdio>
#include<cmath>
#define N 500020
using namespace std;
inline int read(){
    int x=0,f=1;char c;
    do c=getchar(),f=c=='-'?-1:f; while(!isdigit(c));
    do x=(x<<3)+(x<<1)+c-'0',c=getchar(); while(isdigit(c));
    return x*f;
}
struct Complex{
    double x,y;
    Complex(double _=0.0,double __=0.0):x(_),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);}
}a[N],b[N];
typedef long long LL;
const double DFT=2.0,IDFT=-2.0;
const double pi=acos(-1);
LL ans;
int n,len,maxx,Block_Size,Block_num;
int v[N],pos[N],block[N];
int num1[N],num2[N];
inline void FFT(Complex a[],double mode){
    for(int i=0;i<len;i++)
        if(i<pos[i])
            swap(a[i],a[pos[i]]);
    for(int i=2,mid=1;i<=len;i<<=1,mid<<=1){
        Complex wm(cos(2.0*pi/i),sin(mode*pi/i));
        for(int j=0;j<len;j+=i){
            Complex w(1,0);
            for(int k=j;k<j+mid;k++,w=w*wm){
                Complex l=a[k],r=w*a[k+mid];
                a[k]=l+r;a[k+mid]=l-r;
            }
        }
    }
    if(mode==IDFT)
        for(int i=0;i<len;i++)
            a[i].x/=len;
    return;
}
int main(){
    n=read();
    Block_Size=min(n,2000);
    for(int i=1;i<=n;i++){
        v[i]=read(),maxx=max(maxx,v[i]);
        block[i]=(i-1)/Block_Size+1;
        num1[v[i]]++;
    }
    Block_num=block[n];
    for(len=1;len<maxx<<1;len<<=1);
    for(int i=0;i<len;i++){
        pos[i]=pos[i>>1]>>1;
        if(i&1) pos[i]|=len>>1;
    }
    for(int i=0;i<len;i++){
        pos[i]=pos[i>>1]>>1;
        if(i&1) pos[i]|=len>>1;
    }
    for(int i=1;i<=Block_num;i++){
        for(int j=(i-1)*Block_Size+1;j<=min(i*Block_Size,n);j++) num1[v[j]]--;
        for(int j=(i-1)*Block_Size+1;j<=min(i*Block_Size,n);j++){
            for(int k=j+1;k<=min(i*Block_Size,n);k++){
                if(v[j]*2-v[k]>0) ans+=num2[v[j]*2-v[k]];
                if(v[k]*2-v[j]>0) ans+=num1[v[k]*2-v[j]];
            }
            num2[v[j]]++;
        }
    }
    for(int i=2;i<Block_num;i++){
        for(int j=0;j<len;j++)
            a[j]=b[j]=Complex(0,0);
        for(int j=1;j<=(i-1)*Block_Size;j++)
            a[v[j]].x++;
        for(int j=i*Block_Size+1;j<=n;j++)
            b[v[j]].x++;
        FFT(a,DFT);FFT(b,DFT);
        for(int j=0;j<len;j++)
            a[j]=a[j]*b[j];
        FFT(a,IDFT);
        for(int j=(i-1)*Block_Size+1;j<=i*Block_Size;j++)
            ans=ans+(LL)(a[v[j]*2].x+0.1);
    }
    printf("%lld",ans);
    return 0;
}