Codechef Arithmetic Progressions

Discription

Given N integers A₁, A₂, …. A_N, 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 ≤ Nand A_j - A_i = A_k - A_j.

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 A₁, A₂, …, A_N 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)


一种比较容易想到的做法是，我们枚举中间元素，此时两边元素的和就是定值了，这样的话要求的ans就是两边对应的权值数组的卷积中
x^(2*a[j])前面的系数，所以我们跑N次FFT就可以实现了。
然后你发现你TLE了hhh
毕竟这样的复杂度是 O(N*max_element_in_a*log(max_element_in_a))的。

因为你一次FFT只是为了求一项的系数的话，肯定是浪费了很多的时间。

所以我们的想法是用尽量少的次数FFT来得到尽量多的信息。

可以把原序列分成M块，对于每一块，我们计算中间节点在这个块内的答案，此时另外两个点的位置无非以下三种情况：
1.左右端点都在块内：
    这样的话我们枚举中间节点和右端点的位置，而用权值数组O(1)解决左端点。
    总复杂度O(N^2/M)
2.左右端点一个在块内，一个在块外：
    这样就枚举在块内的端点，另一个在块外的端点依然可以用权值数组O(1)解决。
    总复杂度O(N^2/M)
3.左右端点都不在块内：
    这种情况只能枚举中间节点了，左右端点的可能性太多，只能用FFT解决了。
    总复杂度O(N+max_element_in_a*log(max_element_in_a)*M)

当N一定的时候，前两个的复杂度是单减的，最后一个是单增的。
我们解一下两个函数的交点，发现大约M=sqrt(N^2/(max_element_in_a*log(max_element_in_a)))的时候效果最好。
(然而考虑考虑STL的复数运算和FFT的常数的尿性，M还是尽量再小一点比较好)

(结果也的确是这样，反正我分的N/M=sqrt(N*log(N))，速度还可以)

Status
Accepted

Time
1470ms

Memory
20275kB

Length
1935

Lang
C++(gcc 6.3)

Submitted
2018-01-21 16:59:13
Shared

RemoteRunId
17104440
Select Code

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<complex>
#include<algorithm>
#include<cmath>
#include<cstring>
#define ll long long
#define E complex<double>
using namespace std;
#define maxn 70005
const double pi=acos(-1);
ll ans=0;
E a[maxn],b[maxn];
int aa[maxn],bb[maxn];
int n,c[maxn<<1|1],sz,cnt[maxn];
int r[maxn],m,l,bl[100005],mx;

inline void fft(E *d,const ll f){
    for(int i=0;i<m;i++) if(i<r[i]) swap(d[i],d[r[i]]);
    for(int i=1;i<m;i<<=1){
        E omega(cos(pi/i),f*sin(pi/i));
        for(int j=0,p=i<<1;j<m;j+=p){
            E now(1,0);
            for(int k=0;k<i;k++,now*=omega){
                E x=d[j+k],y=d[j+k+i]*now;
                d[j+k]=x+y,d[j+k+i]=x-y;
            }
        }
    }
    
    if(f==-1) for(int i=0;i<m;i++) d[i]/=m;
}

int main(){
    scanf("%d",&n);
    sz=sqrt(n*log(n)/log(2));
    for(int i=1;i<=n;i++){
        scanf("%d",c+i);
        bb[c[i]]++,mx=max(mx,c[i]);
        bl[i]=(i-1)/sz+1;
    }
    
    for(m=1,l=0;m<=(mx<<1);m<<=1) l++;
    for(int i=0;i<m;i++) r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
    
    for(int i=1,j;i<=n;){
        for(j=i;bl[j+1]==bl[i];j++) bb[c[j]]--;
        bb[c[j]]--;
        
        //以下是块内的答案(i和k都是块内)，通过枚举j和k，利用权值数组计算 
        
        for(int k=i;k<j;k++){//枚举中间节点 
            int ba=c[k]<<1;
            for(int u=k+1;u<=j;u++){//枚举右端节点 
                int to=ba-c[u];
                if(to>0&&to<=mx) ans+=(ll)cnt[to];
            }
            cnt[c[k]]++;//记录中间节点左边的块内节点的权值数组 
        }
        for(int k=i;k<j;k++) cnt[c[k]]--;//还原 
        
        //以下是一个块内，一个块外的答案
        
        for(int k=i;k<=j;k++){//枚举中间节点 
            int ba=c[k]<<1;
            for(int u=i;u<k;u++){//枚举块内的左端节点 
                int to=ba-c[u];
                if(to>0&&to<=mx) ans+=(ll)bb[to];                
            }
            for(int u=k+1;u<=j;u++){//枚举块内的右端节点 
                int to=ba-c[u];
                if(to>0&&to<=mx) ans+=(ll)aa[to];
            }
        } 

        //然后利用FFT计算中间节点在本块，两端节点都不是本块的答案 
        for(int k=0;k<=mx;k++){
            a[k]=aa[k],b[k]=bb[k];
        }
        for(int k=mx+1;k<m;k++) a[k]=b[k]=0;
        
        fft(a,1),fft(b,1);
        for(int k=0;k<m;k++) a[k]*=b[k];
        fft(a,-1);
        
        for(int k=i;k<=j;k++) ans+=(ll)(a[c[k]<<1].real()+0.5);
        for(;i<=j;i++) aa[c[i]]++;
    }
    
    printf("%lld\n",ans);
    return 0;
}