[CODECHEF]COUNTARI

Arithmetic Progressions

题目链接：https://www.codechef.com/problems/COUNTARI

All submissions for this problem are available.

Given N integers A1, A2, …. AN, Dexter wants to know how many ways he can choose three numbers such that they are three consecutive terms of an arithmetic progression.
Meaning that, how many triplets (i, j, k) are there such that 1 ≤ i < j < k ≤ N and Aj - Ai = Ak - Aj.
So the triplets (2, 5, 8), (10, 8, 6), (3, 3, 3) are valid as they are three consecutive terms of an arithmetic
progression. But the triplets (2, 5, 7), (10, 6, 8) are not.
Input

First line of the input contains an integer N (3 ≤ N ≤ 100000). Then the following line contains N space separated integers A1, A2, …, AN and they have values between 1 and 30000 (inclusive).
Output

Output the number of ways to choose a triplet such that they are three consecutive terms of an arithmetic progression.
Example

Input:
10
3 5 3 6 3 4 10 4 5 2

Output:
9

Explanation

The followings are all 9 ways to choose a triplet
1 : (i, j, k) = (1, 3, 5), (Ai, Aj, Ak) = (3, 3, 3)
2 : (i, j, k) = (1, 6, 9), (Ai, Aj, Ak) = (3, 4, 5)
3 : (i, j, k) = (1, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)
4 : (i, j, k) = (3, 6, 9), (Ai, Aj, Ak) = (3, 4, 5)
5 : (i, j, k) = (3, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)
6 : (i, j, k) = (4, 6, 10), (Ai, Aj, Ak) = (6, 4, 2)
7 : (i, j, k) = (4, 8, 10), (Ai, Aj, Ak) = (6, 4, 2)
8 : (i, j, k) = (5, 6, 9), (Ai, Aj, Ak) = (3, 4, 5)
9 : (i, j, k) = (5, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)

题意就是给你一个数列，问你这个数列中成等差数列的三元组的个数。
首先我们可以先分块，对于三个全在快内的和只有两个在快内的，我们可以直接暴力去算。
对于只有中间点在块内的，就是枚举两遍的点，这一步可以用FFT优化一下。

$$ans[2*i]=\sum_{d=0}^{i-1}sum[i-d]*sum[i+d]$$
表示的意思就是中间的数为i时的个数。d为枚举的公差。这是一个卷积的形式，直接FFT就好了。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
using namespace std;
#define len 3000
#define MAX 30010
#define LL long long
const int M=100010;
LL ans,pre[M],suc[M],now[M];
int n,num[M],sum,rev[M],dig[M],N,L,Len,maxn;
struct S{
    double x,y;
    void prepare(double xx,double yy){x=xx;y=yy;}
    S operator + (const S&xx){return (S){x+xx.x,y+xx.y};}
    S operator - (const S&xx){return (S){x-xx.x,y-xx.y};}
    S operator * (const S&xx){return (S){x*xx.x-y*xx.y,x*xx.y+y*xx.x};}
}a[M],b[M],c[M],d[M];
inline int in(){
    int x=0; char ch=getchar();
    while(ch<'0'||ch>'9') ch=getchar();
    while(ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    return x;
}
inline void FFT(S *a,int f){
    int i,j,k;
    S wn,w,x,y;
    for(i=0;i<N;++i) d[i]=a[rev[i]];
    for(i=0;i<N;++i) a[i]=d[i];
    for(i=2;i<=N;i<<=1){
        wn.prepare(cos(2*M_PI/i),f*sin(2*M_PI/i));
        for(j=0;j<N;j+=i){
            w.prepare(1,0);
            for(k=j;k<j+i/2;++k){
                x=a[k];
                y=w*a[k+i/2];
                a[k]=x+y;
                a[k+i/2]=x-y;
                w=w*wn;
            }
        }
    }
    if(f==-1) for(i=0;i<N;++i) a[i].x/=(double)N;
}
inline void calc(int x){
    int i,st=(x-1)*len+1,en=min(n,x*len);
    for(i=0;i<N;++i){
        a[i].prepare(pre[i],0);
        b[i].prepare(suc[i],0);
    }
    FFT(a,1);FFT(b,1);
    for(i=0;i<N;++i) c[i]=a[i]*b[i];
    FFT(c,-1);
    for(i=st;i<=en;++i) ans+=(LL)(c[num[i]*2].x+0.5);
}
int main(){
    int i,j,k;
    n=in();
    for(i=1;i<=n;++i){
        num[i]=in();
        suc[num[i]]+=1;
        maxn=max(maxn,num[i]);
    }
    sum=(n/len)+(n%len?1:0);
    for(N=1;N<MAX;N<<=1,L+=1); N<<=1;L+=1;
    for(i=0;i<N;++i){
        for(j=i,Len=0;j;j>>=1) dig[Len++]=j&1;
        for(j=0;j<L;++j) rev[i]=rev[i]*2+dig[j];
    }
    for(i=1;i<=sum;++i){
        int st=(i-1)*len+1,en=min(n,i*len);
        for(j=st;j<=en;++j) suc[num[j]]-=1;
        for(j=st;j<=en;++j){
             for(k=en;k>j;--k){
                int cal=2*num[k]-num[j];
                if(cal>=0&&cal<=maxn) 
                  ans+=now[cal]+suc[cal];
                now[num[k]]+=1;
                cal=2*num[j]-num[k];
                if(cal>=0&&cal<=maxn)
                  ans+=pre[cal]; 
            }
            for(k=j+1;k<=en;++k)
              now[num[k]]-=1;
        }
        calc(i);
        for(j=st;j<=en;++j)
          pre[num[j]]+=1;
    }
    printf("%lld\n",ans);
}