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
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\n",ans);
return 0;
}
By NeighThorn