CodeChef - COUNTARI FTT+分块

Posted 2020-09-30 Meek

Arithmetic Progressions

 

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)

题解：

　　　　考虑分块，分成block块

　　　　假设三个点都在同一块，那么我们就在一块内暴力，复杂度block * ( n/block)  * (n/block)

　　　　假设其中两个点在同一块，那么枚举其中一块的两个点算答案，block * n/block * n/block 

　　·　　假设三个点都不在同一块，枚举中间点属于的那一块 剩下左边和右边进行 FFT， 复杂度block * （n*logn）

　　　　

#include<bits/stdc++.h>
using namespace std;
#pragma comment(linker, "/STACK:102400000,102400000")
#define ls i<<1
#define rs ls | 1
#define mid ((ll+rr)>>1)
#define pii pair<int,int>
#define MP make_pair
typedef long long LL;
typedef unsigned long long ULL;
const long long INF = 1e18+1LL;
const double pi = acos(-1.0);
const int N = 3e5+20, M = 1e6+10, mod = 1e9+7,inf = 2e9;


struct Complex {
    double r , i ;
    Complex () {}
    Complex ( double r , double i ) : r ( r ) , i ( i ) {}
    Complex operator + ( const Complex& t ) const {
        return Complex ( r + t.r , i + t.i ) ;
    }
    Complex operator - ( const Complex& t ) const {
        return Complex ( r - t.r , i - t.i ) ;
    }
    Complex operator * ( const Complex& t ) const {
        return Complex ( r * t.r - i * t.i , r * t.i + i * t.r ) ;
    }
} ;

void FFT ( Complex y[] , int n , int rev ) {
    for ( int i = 1 , j , t , k ; i < n ; ++ i ) {
        for ( j = 0 , t = i , k = n >> 1 ; k ; k >>= 1 , t >>= 1 ) j = j << 1 | t & 1 ;
        if ( i < j ) swap ( y[i] , y[j] ) ;
    }
    for ( int s = 2 , ds = 1 ; s <= n ; ds = s , s <<= 1 ) {
        Complex wn = Complex ( cos ( rev * 2 * pi / s ) , sin ( rev * 2 * pi / s ) ) , w ( 1 , 0 ) , t ;
        for ( int k = 0 ; k < ds ; ++ k , w = w * wn ) {
            for ( int i = k ; i < n ; i += s ) {
                y[i + ds] = y[i] - ( t = w * y[i + ds] ) ;
                y[i] = y[i] + t ;
            }
        }
    }
    if ( rev == -1 ) for ( int i = 0 ; i < n ; ++ i ) y[i].r /= n ;
}
Complex s[N],t[N];

LL cnt[502][30005];
int a[N];
int n,block,pos[N];
LL vis[N];
int main() {
    while(scanf("%d",&n)!=EOF) {
        block = 1500;
        for(int i = 1; i <= n; ++i)
            pos[i] = (i-1)/block + 1;
        int mx = -1;
        for(int i = 0; i <= pos[n]; ++i)
        for(int j = 1; j <= 30000; ++j) cnt[i][j] = 0;
        for(int i = 1; i <= n; ++i) {
            scanf("%d",&a[i]);
            mx = max(mx,a[i]);
            cnt[pos[i]][a[i]]++;
        }

        for(int i = 1; i <= mx; ++i) {
            for(int j = 1; j <= pos[n]; ++j) {
                cnt[j][i] += cnt[j-1][i];
            }
        }
        int len = 1;
        while(len <= 2*mx) len<<=1;
        LL ans = 0;
        for(int k = 1; k <= pos[n]; ++k) {
            for(int i = (k-1)*block + 1; i <= min(k*block,n); ++i) {
                for(int j = i + 1; j <= min(k*block,n); ++j) {
                    if(2*a[i] - a[j] >= 1 && 2*a[i] - a[j] <= mx)
                        ans += cnt[k-1][2*a[i] - a[j]] + vis[2*a[i]-a[j]];
                    if(2*a[j] - a[i] >= 1 &&  2*a[j] - a[i] <= mx)
                        ans += cnt[pos[n]][2*a[j] - a[i]] - cnt[k][2*a[j] - a[i]];
                }
                vis[a[i]] += 1;
            }
            for(int i = (k-1)*block + 1; i <= min(k*block,n); ++i) {
                vis[a[i]] = 0;
            }

            for(int j = 0; j <= mx; ++j)
                s[j] = Complex(cnt[k-1][j],0);
            for(int j = mx+1; j < len; ++j) s[j] = Complex(0,0);

            for(int j = 0; j <= mx; ++j)
                t[j] = Complex(cnt[pos[n]][j] - cnt[k][j] , 0);
            for(int j = mx+1; j < len; ++j) t[j] = Complex(0,0);


            FFT(s,len,1);FFT(t,len,1);
            for(int j = 0; j < len; ++j) s[j] = s[j] * t[j];
            FFT(s,len,-1);

            for(int j = 1; j <= mx; ++j) {
                LL tmp = (LL)(s[2*j].r + 0.5);
                ans += tmp*(cnt[k][j] - cnt[k-1][j]);
            }
        }
        printf("%lld\n",ans);

    }
    return 0;
}