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

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

题面

给出一个长度为n的数组,问有多少三元组\((i,j,k)\)满足\(i<j<k,a_j-a_i=a_k-a_j\)

\(n \leq 10^5, a_i \leq 30000\)

分析

记\(m=\max(a_i)\)

先做一点变形\(a_i+a_k=2a_j\).那么枚举j,把j左边的和右边的所有数字找出来,找出有哪些数字对的和为\(2a_j\)。

维护当前块左边每个数的出现次数cntl,，和右边块的出现次数cntr。

那么数字和为\(2a_j\)的方案数就是\(\sum_{i=0}^{2a_j} cntl(i)cntr(2a_j-i)\)

这是一个卷积的形式,可以FFT.令:

$A(x)=\sum_{i=0}^{m} cntl(i) x^i $

$B(x)=\sum_{i=0}^{m} cntr(i) x^i $

\(A*B\)的第\(2a_j\)项的系数就是答案.但是做n次FFT的复杂度是\(O(n^2 \log n)\),显然是跑不过去的.

我们要平衡一下复杂度，容易想到分块。设块大小为T

i,j,k均不在同一个块的情况

我们枚举j在第几个块,从左到右扫描第2~n-1个块,每块把左边和右边FFT卷积。枚举块内的每个数x,和为\(2x\)的方案数就是卷积结果里第\(2x\)项的系数.

时间复杂度\(O(\frac{n}{B} \times m \log m)\)

i,j,k至少有2个在同一个块的情况

直接暴力从小到大在块内枚举\(k,j(j<k)\).然后讨论j在k前还是在k后。同样维护2个数组cntl表示当前块前面的数的出现次数，cntr表示k后面的数的个数。

如果i在j前,那么答案就是\(cntl[a_j*2-a_k]\)

如果i在k后，那么答案就是\(cntr[a_k*2-a_j]\)

注意到cntl不包含当前块，而cntr包含当前块，这样\(i,j,k\)在同一个块的情况只会被算1次

时间复杂度\(O(\frac{n}{B} \times B \times B)=O(nB)\)

总时间复杂度\(O(nB+\frac{n}{B}m\log m)\)

根据基本不等式,\(B\)取\(\sqrt{m \log m}\)时最优.总时间复杂\(O(n\sqrt{m \log m})\)

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#define maxn 200000
using namespace std;
typedef long double db;
typedef long long ll;
const db pi=acos(-1.0);
struct com{
    double real;
    double imag;
    com(){

    } 
    com(double _real,double _imag){
        real=_real;
        imag=_imag;
    }
    com(double x){
        real=x;
        imag=0;
    }
    void operator = (const com x){
        this->real=x.real;
        this->imag=x.imag;
    }
    void operator = (const double x){
        this->real=x;
        this->imag=0;
    }
    friend com operator + (com p,com q){
        return com(p.real+q.real,p.imag+q.imag);
    }
    friend com operator + (com p,double q){
        return com(p.real+q,p.imag);
    }
    void operator += (com q){
        *this=*this+q;
    }
    void operator += (double q){
        *this=*this+q;
    }
    friend com operator - (com p,com q){
        return com(p.real-q.real,p.imag-q.imag);
    }
    friend com operator - (com p,double q){
        return com(p.real-q,p.imag);
    }
    void operator -= (com q){
        *this=*this-q;
    }
    void operator -= (double q){
        *this=*this-q;
    }
    friend com operator * (com p,com q){
        return com(p.real*q.real-p.imag*q.imag,p.real*q.imag+p.imag*q.real);
    }
    friend com operator * (com p,double q){
        return com(p.real*q,p.imag*q);
    } 
    void operator *= (com q){
        *this=(*this)*q;
    }
    void operator *= (double q){
        *this=(*this)*q;
    }
    friend com operator / (com p,double q){
        return com(p.real/q,p.imag/q);
    } 
    void operator /= (double q){
        *this=(*this)/q;
    } 
    void print(){
        printf("%lf + %lf i ",real,imag);
    }
};
void fft(com *x,int n,int type){
	static int rev[maxn+5];
	int dn=1,k=0;
	while(dn<n){
		dn*=2;
		k++;
	}
	for(int i=0;i<n;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(k-1));
	for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
	for(int len=1;len<n;len*=2){
		int sz=len*2;
		com wn1=com(cos(2*pi/sz),sin(2*pi/sz)*type);
		for(int l=0;l<n;l+=sz){
			int r=l+len-1;
			com wnk=1;
			for(int i=l;i<=r;i++){
				com tmp=x[i+len];
				x[i+len]=x[i]-wnk*tmp;
				x[i]=x[i]+wnk*tmp;
				wnk*=wn1;
			}
		}
	}
	if(type==-1) for(int i=0;i<n;i++) x[i]/=n;
}
void mul(com *a,com *b,com *ans,int n){
	fft(a,n,1);
	if(a!=b) fft(b,n,1);
	for(int i=0;i<n;i++) ans[i]=a[i]*b[i];
	fft(ans,n,-1);
}

int n;
int bsz,bcnt;
int a[maxn+5]; 
int bel[maxn+5];
int lb[maxn+5],rb[maxn+5];//预处理每个数的左右边界，卡常 
void ini_block(){
	bsz=sqrt(n*log(n)/log(2));
	bcnt=1;
	for(int i=1;i<=n;i++){
		bel[i]=bcnt;
		if(i%bsz==0) bcnt++;
	}
	for(int i=1;i<=bcnt;i++){
		lb[i]=bsz*(i-1)+1;
		rb[i]=min(i*bsz,n); 
	}
} 

int cntl[maxn+5],cntr[maxn+5];
com fl[maxn+5],fr[maxn+5],res[maxn+5];
int main(){
	scanf("%d",&n);
	ini_block();
	int maxv=0;
	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
		maxv=max(maxv,a[i]);
	} 
	//i,j,k有>=2个在同一块内，直接暴力O(n/B*B^2)=O(Bn)
	ll ans=0;
	for(int i=n;i>=1;i--) cntr[a[i]]++;
	for(int id=1;id<=bcnt;id++){
		//cntl记录的是当前块前面的数的出现情况
		//cntr记录的是i后面的数的出现情况
		//注意i,j,k均在一个块内的情况只能被算1次  
		for(int k=lb[id];k<=rb[id];k++){
			cntr[a[k]]--;
			for(int j=lb[id];j<k;j++){
				int ai=a[j]*2-a[k];
				if(ai>=0&&ai<=maxn) ans+=cntl[ai]; 
				ai=a[k]*2-a[j];
				if(ai>=0&&ai<=maxn) ans+=cntr[ai];//这样只有这个时候会算到i,j,k在同一个块里的情况 
			}
			
		}
		for(int j=lb[id];j<=rb[id];j++) cntl[a[j]]++;
	} 
	memset(cntl,0,sizeof(cntl));
	memset(cntr,0,sizeof(cntr));
	for(int i=n;i>=lb[2];i--) cntr[a[i]]++;
	for(int i=1;i<=rb[1];i++) cntl[a[i]]++; 
	for(int id=2;id<bcnt;id++){
		for(int j=lb[id];j<=rb[id];j++) cntr[a[j]]--;
		for(int j=0;j<=maxv;j++){
			fl[j]=cntl[j];
			fr[j]=cntr[j];
		}
		int dn=1;
		while(dn<=maxv*2) dn*=2;
		mul(fl,fr,res,dn);
		for(int j=lb[id];j<=rb[id];j++){
			ans+=(ll)(res[a[j]*2].real+0.5);
		}
		for(int j=0;j<=dn;j++){
			fl[j]=0;
			fr[j]=0;
			res[j]=0;
		}
		for(int j=lb[id];j<=rb[id];j++) cntl[a[j]]++;
	} 
	printf("%lld\n",ans);
}