CC Arithmetic Progressions (FFT + 分塊處理)

轉載請註明出處，謝謝http://blog.csdn.net/ACM_cxlove?viewmode=contents    by---cxlove

題目：給出n個數，選出三個數，按下標順序形成等差數列

http://www.codechef.com/problems/COUNTARI

如果只是形成 等差數列並不難，大概就是先求一次卷積，然後再O(n)枚舉，判斷2 * a[i]的種數，不過按照下標就不會了。

有種很矬的，大概就是O(n)枚舉中間的數，然後 對兩邊分別卷積，O(n * n * lgn)。

如果能想到枚舉中間的數的話，應該可以進一步想到分塊處理。

如果分為K塊

那麼分為幾種情況 ：

三個數都是在當前塊中，那麼可以枚舉後兩個數，尋找第一個數，複雜度O(N/K * N/K)

兩個數在當前塊中，那麼另外一個數可能在前面，也可能在後面，同理還是枚舉兩個數，尋找，複雜度
O(N/K * N/K)

如果只有一個數在當前塊中，那麼就要對兩邊的數進行卷積，然後枚舉當前塊中的數，查詢2 × a[i]。複雜度O(N * lg N)

那麼總體就是O(k * (N/K * N/K + N * lg N))。

#include <iostream>#include <cstdio>#include <cstring>#include <cmath>#include <algorithm>using namespace std;//FFT copy from kuangbinconst double pi = acos (-1.0);// Complex  z = a + b * i  struct Complex {    double a, b;    Complex(double _a=0.0,double _b=0.0):a(_a),b(_b){}    Complex operator + (const Complex &c) const {        return Complex(a + c.a , b + c.b);    }    Complex operator - (const Complex &c) const {        return Complex(a - c.a , b - c.b);    }    Complex operator * (const Complex &c) const {        return Complex(a * c.a - b * c.b , a * c.b + b * c.a);    }};//len = 2 ^ kvoid change (Complex y[] , int len) {    for (int i = 1 , j = len / 2 ; i < len -1 ; i ++) {        if (i < j) swap(y[i] , y[j]);        int k = len / 2;        while (j >= k) {            j -= k;            k /= 2;        }        if(j < k) j += k;    } }// FFT // len = 2 ^ k// on = 1  DFT    on = -1 IDFTvoid FFT (Complex y[], int len , int on) {    change (y , len);    for (int h = 2 ; h <= len ; h <<= 1) {        Complex wn(cos (-on * 2 * pi / h), sin (-on * 2 * pi / h));        for (int j = 0 ; j < len ; j += h) {            Complex w(1 , 0);            for (int k = j ; k < j + h / 2 ; k ++) {                Complex u = y[k];                Complex t = w * y [k + h / 2];                y[k] = u + t;                y[k + h / 2] = u - t;                w = w * wn;            }        }    }    if (on == -1) {        for (int i = 0 ; i < len ; i ++) {            y[i].a /= len;        }    }}const int N = 100005;typedef long long LL;int n , a[N];int block , size;LL num[N << 2];int min_num = 30000 , max_num = 1;int before[N] = {0}, behind[N] = {0} , in[N] = {0};Complex x1[N << 2] ,x2[N << 2];int main () {    #ifndef ONLINE_JUDGE        freopen("input.txt" , "r" , stdin);    #endif    scanf ("%d", &n);    for (int i = 0 ; i < n ; ++ i) {        scanf ("%d", &a[i]);        behind[a[i]] ++;        min_num = min (min_num , a[i]);        max_num = max (max_num , a[i]);    }    LL ret = 0;    block = min(n , 35);    size = (n + block - 1) / block;    for (int t = 0 ; t < block ; t ++) {        int s = t * size , e = (t + 1) * size;        if (e > n) e = n;        for (int i = s ; i < e ; i ++) {            behind[a[i]] --;        }        for (int i = s ; i < e ; i ++) {            for (int j = i + 1 ; j < e ; j ++) {                int m = 2 * a[i] - a[j];                if(m >= 1 && m <= 30000) {                     // both of three in the block                    ret += in[m];                    // one of the number in the pre block                    ret += before[m];                }                m = 2 * a[j] - a[i];                if (m >= 1 && m <= 30000) {                    // one of the number in the next block                    ret += behind[m];                }            }            in[a[i]] ++;        }        // pre block , current block , next block        if (t > 0 && t < block - 1) {            int l = 1;            int len = max_num + 1;            while (l < len * 2) l <<= 1;            for (int i = 0 ; i < len ; i ++) {                x1[i] = Complex (before[i] , 0);            }            for (int i = len ; i < l ; i ++) {                x1[i] = Complex (0 , 0);            }            for (int i = 0 ; i < len ; i ++) {                x2[i] = Complex (behind[i] , 0);            }            for (int i = len ; i < l ; i ++) {                x2[i] = Complex (0 , 0);            }            FFT (x1 , l , 1);            FFT (x2 , l , 1);            for (int i = 0 ; i < l ; i ++) {                x1[i] = x1[i] * x2[i];            }            FFT (x1 , l , -1);            for (int i = 0 ; i < l ; i ++) {                num[i] = (LL)(x1[i].a + 0.5);            }            for (int i = s ; i < e ; i ++) {                ret += num[a[i] << 1];            }        }        for (int i = s ; i < e ; i ++) {            in[a[i]] --;            before[a[i]] ++;        }    }    printf("%lld\n", ret);    return 0;}