bzoj 3509 [CodeChef] COUNTARI

3509: [CodeChef] COUNTARI

Time Limit: 40 Sec   Memory Limit: 128 MB
Submit: 1022   Solved: 300
[ 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

【分析】

分块暴力FFT。

迷之复杂度。竟然没被卡掉。

对于每个点j，朴素算法是扫一遍左边，扫一遍右边，构造指数生成函数，FFT跑一遍统计答案。

我们发现复杂度O(n*maxv*log(maxv))，貌似并没有n^2扫一遍好。

每次做一遍FFT太慢了，我们想一遍FFT统计多个点的答案。

于是分一个块...块外FFT，块内扫一遍统计= =但是复杂度&*^#$@&$!，总之不会卡掉哒

【代码】

#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#define pi acos(-1)
#define ll long long
#define M(a) memset(a,0,sizeof a)
#define fo(i,j,k) for(i=j;i<=k;i++)
using namespace std;
const int mxn=120005;
ll ans;
int n,m,T,N,M,L,size,num;
struct kuai {int l,r;} t[mxn];
int w[mxn],R[mxn],lc[mxn],rc[mxn],tmp[mxn];
struct E
{
	double r,f;
	E (double u,double v) {r=u,f=v;} 
	E () {}
	E operator + (E u) {return E(r+u.r,f+u.f);} 
	E operator - (E u) {return E(r-u.r,f-u.f);}
	E operator * (E u) {return E(r*u.r-f*u.f,r*u.f+f*u.r);}
	E operator / (int u) {return E(r/u,f/u);}
}a[mxn],b[mxn],c[mxn];
inline void FFT(E *a,int f)
{
	int i,j,k;
	fo(i,0,n-1) if(i<R[i]) swap(a[i],a[R[i]]);
	for(i=1;i<n;i<<=1)
	{
		E wn(cos(pi/i),f*sin(pi/i));
		for(j=0;j<n;j+=(i<<1))
		{
			E w(1,0);
			for(k=0;k<i;k++,w=w*wn)
			{
				E x=a[j+k],y=w*a[j+k+i];
				a[j+k]=x+y,a[j+k+i]=x-y;
			}
		}
	}
	if(f==-1) fo(i,0,n-1) a[i]=a[i]/n;
}
inline void solve(int p)
{
	int i,j,k;M(a),M(b),M(tmp);
	fo(i,1,t[p].l-1) a[w[i]].r+=1;
	fo(i,t[p].r+1,N) b[w[i]].r+=1;
	FFT(a,1),FFT(b,1);
	fo(i,0,n-1) a[i]=a[i]*b[i];
	FFT(a,-1);
	fo(i,t[p].l,t[p].r) ans+=(int)(a[w[i]*2].r+0.5),tmp[w[i]]++;
	fo(i,t[p].l,t[p].r)
	{
		tmp[w[i]]--;
		fo(j,t[p].l,i-1)
		  if((k=w[i]+w[i]-w[j])>=0)
		    ans-=tmp[k];
	}
	fo(i,t[p].l,t[p].r)
	{
		rc[w[i]]--;
		fo(j,t[p].l,i-1)
		  if((k=w[i]+w[i]-w[j])>=0)
		    ans+=rc[k];
		fo(j,i+1,t[p].r)
		  if((k=w[i]+w[i]-w[j])>=0)
		    ans+=lc[k];
		lc[w[i]]++;
	}
}
int main()
{
	int i,j;
	scanf("%d",&N);
	fo(i,1,N) scanf("%d",&w[i]);
	size=1500;
	fo(i,1,N)
	{
		t[i].l=(i-1)*size+1;
		t[i].r=i*size;
		if(t[i].r>=N)
		{
			t[i].r=N,num=i;  //块的数量 
			break;
		}
	}
	m=60000;for(n=1;n<=m;n<<=1) L++;
	fo(i,0,n-1) R[i]=(R[i>>1]>>1)|((i&1)<<L-1);
	fo(i,1,N) rc[w[i]]++;
	fo(i,1,num) solve(i);
	printf("%lld\n",ans);
	return 0;
}