【BNU】40719 Arithmetic Progressions【分块+FFT】

传送门：【BNU】40719 Arithmetic Progressions

题目分析：

用分块+FFT强行AC了这题……
之前一直TLE……然后改了好久把姿势改的优美点了……终于过了……

大概思路是：我们考虑分块，假设每一块的大小为S，一共分了B块然后我们分两种情况讨论：
1.第二个数在第i块，第一个数在(1~i-1)块内，第三个数在(i+1~B)块内。
2.至少两个数在同一块内。

对于第一种情况，我们可以用$FFT$实现，每块的复杂度为$O(M\log M)$，$M$为数的大小，这题大概$M=2^{16}$。

对于第二种情况，我们可以$O((\frac{N}{S})^2)$的复杂度内实现。

具体实现就不多说了，知道大体该往哪个方向思考就好了。

总复杂度：$O(SM\log M + \frac{N^2}{S})$

我写的常数大的不行啊= =不知道正解到底是怎么样，反正比较无脑的就是这样暴力……

$my~~code:$

#include <stdio.h>
#include <string.h>
#include <math.h>
#include <algorithm>
using namespace std ;

typedef long long LL ;

#define clr( a , x ) memset ( a , x , sizeof a )
#define cpy( a , x ) memcpy ( a , x , sizeof a )

const int MAXN = 100005 ;
const int SQR = 330 ;
const double pi = acos ( -1.0 ) ;

struct Complex {
    double r , i ;
    Complex ( double r = 0 , double i = 0 ) : 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 ) ;
    }
} ;

int n ;
Complex x1[MAXN << 2] ;
Complex x2[MAXN << 2] ;
int cnt[MAXN] ;
int a[MAXN] ;
int in[MAXN] , L[MAXN] ;

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

void solve () {
    LL ans = 0 ;
    int N = 1 << 16 ;
    int sqr = n / 50 + 1 ;
    clr ( cnt , 0 ) ;
    clr ( L , 0 ) ;
    clr ( in , 0 ) ;
    for ( int i = 0 ; i < n ; ++ i ) {
        scanf ( "%d" , &a[i] ) ;
        ++ cnt[a[i]] ;
    }
    for ( int l = 0 ; l < n ; l += sqr ) {
        int r = min ( n , l + sqr ) ;
        for ( int i = l ; i < r ; ++ i ) {
            ++ in[a[i]] ;
        }
        for ( int i = 0 ; i < N ; ++ i ) {
            x1[i] = Complex ( L[i] , 0 ) ;
        }
        for ( int i = 0 ; i < N ; ++ i ) {
            x2[i] = Complex ( cnt[i] - in[i] - L[i] , 0 ) ;
        }
        FFT ( x1 , N , 1 ) ;
        FFT ( x2 , N , 1 ) ;
        for ( int i = 0 ; i < N ; ++ i ) {
            x1[i] = x1[i] * x2[i] ;
        }
        FFT ( x1 , N , -1 ) ;
        for ( int i = l ; i < r ; ++ i ) {
            if ( a[i] * 2 < N ) ans += ( LL ) ( x1[a[i] << 1].r + 0.5 ) ;
        }
        for ( int i = l ; i < r ; ++ i ) {
            -- in[a[i]] ;
        }
        for ( int i = l ; i < r ; ++ i ) {
            for ( int j = l ; j <= i ; ++ j ) {
                ++ in[a[j]] ;
            }
            for ( int j = i + 1 ; j < r ; ++ j ) {
                in[a[j]] ++ ;
                int x = a[i] * 2 - a[j] ;
                int y = a[j] * 2 - a[i] ;
                if ( x >= 0 ) ans += L[x] ;
                if ( y >= 0 ) ans += cnt[y] - in[y] - L[y] ;
            }
            for ( int j = l ; j < r ; ++ j ) {
                -- in[a[j]] ;
            }
        }
        for ( int i = l ; i < r ; ++ i ) {
            ++ L[a[i]] ;
        }
    }
    printf ( "%lld\n" , ans ) ;
}

int main () {
    while ( ~scanf ( "%d" , &n ) ) solve () ;
    return 0 ;
}