CodeChef Arithmetic Progressions (分块FFT)

题意:给出 A 1 ,A 2 ,...,A N 统计满足:
▶ i < j < k
▶ A i + A k = 2A j
的 (i,j,k) 数量。

(N ≤ 30000,A i ≤ 10 5 )

分块,把长度为n的序列分成sqrt(n)块长度为sqrt(n)的序列,然后遍历每一个块分三种情况:

1)三个都在同一个块里面:

暴力枚举后两个,每次维护前面的数的个数,复杂度O(sqrt(n)*n)

2)两个在同一块里面:

暴力枚举块中的两个,维护块前的数的个数和块后的数的个数,复杂度O(sqrt(n)*n)

3)一个在块中,一个在块前,一个在块后

块前块后做fft,枚举块中的数,复杂度O(nlgn*sqrt(n)).

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
using namespace std;
#define pi acos (-1)
#define maxn 151111

struct plex {
    double x, y;
    plex (double _x = 0.0, double _y = 0.0) : x (_x), y (_y) {}
    plex operator + (const plex &a) const {
        return plex (x+a.x, y+a.y);
    }
    plex operator - (const plex &a) const {
        return plex (x-a.x, y-a.y);
    }
    plex operator * (const plex &a) const {
        return plex (x*a.x-y*a.y, x*a.y+y*a.x);
    }
};

void change (plex 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;
    }
}

void fft(plex y[],int len,int on)
{
    change(y,len);
    for(int h = 2; h <= len; h <<= 1)
    {
        plex wn(cos(on*2*pi/h),sin(on*2*pi/h));
        for(int j = 0;j < len;j+=h)
        {
            plex w(1,0);
            for(int k = j;k < j+h/2;k++)
            {
                plex u = y[k];
                plex 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].x /= len;
}
int num[maxn], num2[maxn];
long long sum[maxn];
int a[maxn];
plex x[maxn], y[maxn];
int n;

int main () {
    //freopen ("in.txt", "r", stdin);
    scanf ("%d", &n);
    for (int i = 0; i < n; i++) {
        scanf ("%d", &a[i]);
    }
    long long ans = 0;
    int block = sqrt (n), len = n/block;
    if (len*block != n)
        block++;
    //1 都在当前块
    for (int t = 0; t < block; t++) {
        memset (num, 0, sizeof num);
        for (int i = min ((t+1)*len-1, n-1); i >= t*len; i--) {
            for (int j = t*len; j < i; j++)
                num[a[j]]++;
            for (int j = i-1; j >= t*len; j--) {
                num[a[j]]--;
                int cur = 2*a[j]-a[i];
                if (cur > 0)
                    ans += num[cur];
            }
        }
    }
    //两个在当前块
    for (int t = 0; t < block; t++) {
        memset (num, 0, sizeof num);
        //第三个在前面的块中
        for (int i = 0; i < t*len; i++)
            num[a[i]]++;
        for (int i = t*len; i < n && i < (t+1)*len; i++) {
            for (int j = i+1; j < n && j < (t+1)*len; j++) {
                int cur = a[i]*2-a[j];
                if (cur > 0)
                    ans += num[cur];
            }
        }
        //第三个在后面的块中
        memset (num, 0, sizeof num);
        for (int i = (t+1)*len; i < n; i++)
            num[a[i]]++;
        for (int i = t*len; i < n && i < (t+1)*len; i++) {
            for (int j = i+1; j < n && j < (t+1)*len; j++) {
                int cur = a[j]*2-a[i];
                if (cur > 0)
                    ans += num[cur];
            }
        }
    }
    //只有一个在当前块中 一个在前面 一个在后面
    for (int t = 0; t < block; t++) {
        int cnt1 = 0, cnt2 = 0, Max = 0;
        memset (num, 0, sizeof num);
        for (int i = 0; i < t*len; i++) {
            num[a[i]]++;
            Max = max (Max, a[i]);
        }
        memset (num2, 0, sizeof num2);
        for (int i = (t+1)*len; i < n; i++) {
            num2[a[i]]++;
            Max = max (Max, a[i]);
        }
        int l = 1;
        Max++;
        while (l < 2*Max)
            l <<= 1;
        for (int i = 0; i < l; i++) {
            x[i] = plex (num[i], 0);
        }
        for (int i = 0; i < l; i++) {
            y[i] = plex (num2[i], 0);
        }
        fft (x, l, 1);
        fft (y, l, 1);
        for (int i = 0; i < l; i++)
            x[i] = x[i]*y[i];
        fft (x, l, -1);
        memset (sum, 0, sizeof sum);
        for (int i = 1; i < l; i++) {
            sum[i] = (long long) (x[i].x+0.5);
        }
        for (int i = t*len; i < (t+1)*len && i < n; i++) {
            ans += sum[2*a[i]];
        }
    }
    printf ("%lld\n", ans);
    return 0;
}