题目链接:http://www.codechef.com/problems/COUNTARI/
题意:给出一个长度为n的数列A。求有多少三元组(i,j,k)满足i<j<k且A[i]-A[j]=A[j]-A[k]?
思路:将数列分成K块。那么对于每三个数:
(1)在一块中,枚举两个,查找第三个数的个数;
(2)两个数在一块中,另一个数在前面或者后面的块中,也是枚举该块内的两个数,查找第三个数的个数;
(3)三个数在三块中,枚举中间一块中的数,用FFT计算前后两部分。
struct node
{
double x,y;
node(double _x=0.0,double _y=0.0)
{
x=_x;
y=_y;
}
node operator+(node a)
{
return node(x+a.x,y+a.y);
}
node operator-(node a)
{
return node(x-a.x,y-a.y);
}
node operator*(node a)
{
return node(x*a.x-y*a.y,x*a.y+y*a.x);
}
};
node A[N];
int L;
int reverse(int x)
{
int ans=0,i;
FOR0(i,L) if(x&(1<<i)) ans|=1<<(L-1-i);
return ans;
}
void bitReverseCopy(node a[],int n)
{
int i;
FOR0(i,n) A[i]=a[i];
FOR0(i,n)
{
a[reverse(i)]=A[i];
}
}
void fft(node a[],int n,int on)
{
bitReverseCopy(a,n);
int len,i,j,k;
node x,y,u,t;
for(len=2;len<=n;len<<=1)
{
x=node(cos(-on*2*PI/len),sin(-on*2*PI/len));
for(j=0;j<n;j+=len)
{
y=node(1,0);
for(k=j;k<j+len/2;k++)
{
u=a[k];
t=y*a[k+len/2];
a[k]=u+t;
a[k+len/2]=u-t;
y=y*x;
}
}
}
if(on==-1)
{
FOR0(i,n) a[i].x/=n;
}
}
int a[N],C[N],n;
int size;
i64 ans;
int b[100][2],bNum;
void deal1()
{
int L,R,i,j,k;
for(L=1;L<=n;L+=size)
{
R=min(n,L+size-1);
b[bNum][0]=L;
b[bNum][1]=R;
bNum++;
for(i=L;i<=R;i++)
{
for(j=i+1;j<=R;j++)
{
k=2*a[i]-a[j];
if(k>=0) ans+=C[k];
}
C[a[i]]++;
}
for(i=L;i<=R;i++) C[a[i]]--;
}
}
void deal2()
{
int L,R,i,j,k,t;
for(t=0;t<bNum;t++)
{
L=b[t][0];
R=b[t][1];
for(i=L;i<=R;i++) for(j=i+1;j<=R;j++)
{
k=2*a[i]-a[j];
if(k>=0) ans+=C[k];
}
for(i=L;i<=R;i++) C[a[i]]++;
}
for(i=1;i<=n;i++) C[a[i]]--;
for(t=bNum-1;t>=0;t--)
{
L=b[t][0];
R=b[t][1];
for(i=R;i>=L;i--) for(j=i-1;j>=L;j--)
{
k=2*a[i]-a[j];
if(k>=0) ans+=C[k];
}
for(i=L;i<=R;i++) C[a[i]]++;
}
for(i=1;i<=n;i++) C[a[i]]--;
}
int pre[N],tail[N];
node P[N],Q[N];
int M;
i64 Ans[N];
void init()
{
M=1; L=0;
while(M<=60000) M<<=1,L++;
int i;
for(i=0;i<=30000;i++) P[i]=node(pre[i],0),Q[i]=node(tail[i],0);
while(i<M) P[i]=node(0,0),Q[i]=node(0,0),i++;
fft(P,M,1);
fft(Q,M,1);
for(i=0;i<M;i++) P[i]=P[i]*Q[i];
fft(P,M,-1);
for(i=0;i<M;i++) Ans[i]=(i64)(P[i].x+0.5);
}
void deal3()
{
int i,j,k,t,L,R;
L=b[0][0]; R=b[0][1];
for(i=L;i<=R;i++) pre[a[i]]++;
for(t=1;t<bNum;t++)
{
L=b[t][0];
R=b[t][1];
for(i=L;i<=R;i++) tail[a[i]]++;
}
for(t=1;t<=bNum-2;t++)
{
L=b[t][0];
R=b[t][1];
for(i=L;i<=R;i++) tail[a[i]]--;
init();
for(i=L;i<=R;i++) ans+=Ans[a[i]<<1];
for(i=L;i<=R;i++) pre[a[i]]++;
}
}
int main()
{
RD(n);
int i;
FOR1(i,n) RD(a[i]);
size=n/min(n,35); ans=0;
deal1();
deal2();
deal3();
printf("%lld\n",ans);
}