BZOJ 3509: [CodeChef] COUNTARI(fft+分块)

题面

Description

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

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

一行一个整数。
Sample Input
10

3 5 3 6 3 4 10 4 5 2

Sample Output
9

解题思路

　　比较有意思的一道题。首先直接的思路就是对于一个点来说，把它两边卷起来算贡献，但这样复杂度是(O(n^2logn))的，无法接受。考虑分块，我们其实可以把原序列分块，然后每次(fft)块两边，这样算出来的贡献块内都可以用。对于块内的贡献可以直接暴力预处理出，预处理就是用两个数组，一个表示左边一个表示右边，然后先枚举块，再枚举两重块内元素。设块的大小为(T)，那么就有(n/T)块，时间复杂度为(O(n*T+dfrac{n}{T}*sqrt{n*logn}))，然后两边平衡一下发现(T)取(sqrt{nlogn})时最优。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<cmath>
#include<algorithm>

using namespace std;
const int MAXN = 100005;
const double Pi = acos(-1);
typedef long long LL;

inline int rd(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch)) {f=ch=='-'?0:1;ch=getchar();}
	while(isdigit(ch))  {x=(x<<1)+(x<<3)+ch-'0';ch=getchar();}
	return f?x:-x;
}

int n,a[MAXN],bl[MAXN],l[MAXN],r[MAXN];
int siz,num,rev[MAXN<<2],limit=1,mx;
int L[30005],R[30005];
LL ans;
struct Complex{
	double x,y;
	Complex(double xx=0,double yy=0){
		x=xx;y=yy;
	}
}f[30005<<2],g[30005<<2];

Complex operator+(const Complex A,const Complex B){return Complex(A.x+B.x,A.y+B.y);}
Complex operator-(const Complex A,const Complex B){return Complex(A.x-B.x,A.y-B.y);}
Complex operator*(const Complex A,const Complex B){return Complex(A.x*B.x-A.y*B.y,A.x*B.y+A.y*B.x);}

inline void fft(Complex *f,int type){
	for(int i=0;i<limit;i++)
		if(i<rev[i]) swap(f[i],f[rev[i]]);
	Complex Wn,w,tmp;int len;
	for(int p=2;p<=limit;p<<=1){
		len=p>>1;Wn=Complex(cos(Pi/len),type*sin(Pi/len));
		for(int k=0;k<limit;k+=p){
			w=Complex(1,0);
			for(int t=k;t<k+len;t++){
				tmp=f[t+len]*w;f[t+len]=f[t]-tmp;
				f[t]=f[t]+tmp;w=w*Wn;
			}
		}
	}
	if(type==-1) for(int i=0;i<limit;i++) f[i].x/=limit;
}

inline int max(int x,int y){
	return x>y?x:y;
}

signed main(){
	n=rd();siz=sqrt(n*log2(n))+1;num=n/siz;if(n%siz) siz++; 
	for(int i=1;i<=n;i++) {a[i]=rd();bl[i]=(i-1)/siz+1;R[a[i]]++;}
	for(int i=1;i<=num;i++) l[i]=(i-1)*siz+1,r[i]=i*siz;r[num]=n;int tmp;
	for(int i=1;i<=num;i++){
		for(int j=l[i];j<=r[i];j++) R[a[j]]--;
		for(int j=l[i];j<=r[i];j++){
			for(int k=j+1;k<=r[i];k++){
				tmp=a[j]*2-a[k];
				if(tmp>=0 && tmp<=30000) ans+=L[tmp];
				tmp=a[k]*2-a[j];
				if(tmp>=0 && tmp<=30000) ans+=R[tmp];
			}
			L[a[j]]++;
		}
	}memset(L,0,sizeof(L));
	for(int i=2;i<num;i++){
		mx=-1;limit=1;
		for(int j=1;j<l[i];j++) mx=max(mx,a[j]),L[a[j]]++;
		for(int j=r[i]+1;j<=n;j++) mx=max(mx,a[j]),R[a[j]]++;
		for(int j=0;j<=mx;j++) f[j].x=L[j],g[j].x=R[j];
		while(limit<=mx*2) limit<<=1;
		for(int j=0;j<limit;j++) rev[j]=(rev[j>>1]>>1)|((j&1)?limit>>1:0);
		fft(f,1);fft(g,1);for(int j=0;j<limit;j++) f[j]=f[j]*g[j];fft(f,-1);
		for(int j=l[i];j<=r[i];j++) ans+=(LL)(f[a[j]<<1].x+0.5);
		for(int j=0;j<=mx;j++) L[j]=R[j]=0;
		for(int j=0;j<limit;j++) f[j].x=f[j].y=g[j].x=g[j].y=0;
	}
	printf("%lld
",ans);
	return 0;
}