BZOJ 3509: [CodeChef] COUNTARI

3509: [CodeChef] COUNTARI

Time Limit: 40 Sec  Memory Limit: 128 MB
Submit: 883  Solved: 250
[Submit][Status][Discuss]

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

HINT

Source

分析：

考虑枚举中间项$a[j]$，我们可以使用$FFT$快速计算$ik$的配对个数，但是这样显然是$O(Nmaxnumlgmaxnum)$的，还没有暴力跑的快...

所以可以考虑分块，对于$ijk$都在同一块内或者$ij$在一块或者$jk$在一块的情况我们预处理出来，这些的复杂度是$O(len*N)$的，然后$ik$位于和$j$不同快的情况我们用$FFT$来计算，这个的复杂度是$O(frac {N}{len}*maxnumlgmaxnum)$的...

如此计算，块的大小大概开到3000~2000就可以过去了...

代码：

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
using namespace std;

const int maxn=100000+5,maxblock=50+5;
const double pi=acos(-1);

int n,N,m,L,blo,a[maxn],R[maxn],id[maxn],be[maxblock],en[maxblock],cnt[maxblock][maxn];
long long ans;

struct complex{
	
	double r,i;
	
	inline complex(double a=0,double b=0): r(a),i(b) {};
	
	inline complex operator + (const complex &a){
		return complex(r+a.r,i+a.i);
	}
	
	inline complex operator - (const complex &a){
		return complex(r-a.r,i-a.i);
	}
	
	inline complex operator * (const complex &a){
		return complex(r*a.r-i*a.i,r*a.i+i*a.r);
	}
	
}l[maxn],s[maxn];

inline void FFT(complex *l,int f){
	for(int i=0;i<N;i++)
		if(i>R[i]) swap(l[i],l[R[i]]);
	for(int i=1;i<N;i<<=1){
		complex wn(cos(pi/i),f*sin(pi/i));
		for(int j=0;j<N;j+=i<<1){
			complex w(1,0);
			for(int k=0;k<i;k++,w=w*wn){
				complex x=l[j+k],y=l[j+k+i]*w;
				l[j+k]=x+y;
				l[j+k+i]=x-y;
			}
		}
	}
	if(f==-1)
		for(int i=0;i<N;i++)
			l[i].r/=N;
}

signed main(void){
//	freopen("countari.in","r",stdin);
//	freopen("countari.out","w",stdout);
	scanf("%d",&n);
	L=0;ans=0;
	for(int i=1;i<=n;i++)
		scanf("%d",&a[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));
	blo=min(int(sqrt(n)*10),n);
	for(int i=1;i<=n;i++) id[i]=(i-1)/blo+1;
	for(int i=1;i<=id[n];i++)
		be[i]=lower_bound(id+1,id+n+1,i)-id,
		en[i]=upper_bound(id+1,id+n+1,i)-id-1;
	memset(cnt,0,sizeof(cnt));
	for(int i=1;i<=n;i++) cnt[id[i]][a[i]]++;
	for(int B=1;B<=id[n];B++){
		for(int i=be[B];i<=en[B];i++){
			cnt[B][a[i]]--;
			for(int j=be[B];j<i;j++)
				if(2*a[i]-a[j]>=0) ans+=cnt[B][2*a[i]-a[j]];
		}
	}
	memset(cnt,0,sizeof(cnt));
	for(int i=1;i<=n;i++) cnt[0][a[i]]++;
	for(int B=1;B<=id[n];B++){
		for(int i=be[B];i<=en[B];i++) cnt[0][a[i]]--;
		for(int i=be[B];i<=en[B];i++){
			for(int j=be[B];j<i;j++)
				if(2*a[i]-a[j]>=0) ans+=cnt[0][2*a[i]-a[j]];
		}
	}
	memset(cnt,0,sizeof(cnt));
	for(int i=1;i<=n;i++) cnt[0][a[i]]++;
	for(int B=id[n];B>=1;B--){
		for(int i=be[B];i<=en[B];i++) cnt[0][a[i]]--;
		for(int i=be[B];i<=en[B];i++){
			for(int j=i-1;j>=be[B];j--)
				if(2*a[j]-a[i]>=0) ans+=cnt[0][2*a[j]-a[i]];
		}
	}
	for(int B=1;B<=id[n];B++){
		for(int i=0;i<N;i++) l[i].r=l[i].i=s[i].r=s[i].i=0;
		for(int i=1;i<be[B];i++) l[a[i]].r++;
		for(int i=en[B]+1;i<=n;i++) s[a[i]].r++;
		FFT(l,1);FFT(s,1);
		for(int i=0;i<N;i++) l[i]=l[i]*s[i];
		FFT(l,-1);
		for(int i=be[B];i<=en[B];i++) ans+=int(l[a[i]*2].r+0.5);
	}
	printf("%lld
",ans);
	return 0;
}

　　

---

By NeighThorn