CodeChef COUNTARI Arithmetic Progressions FFT + 分块

题目大意：

就是现在给出一个数列A[1~n], 每个数都是不超过30000的正整数, 现在求有多少个三元组(i, j, k)满足 1 <= i < j < k <= n使得 A[i], A[j], A[k]成等差数列 n <= 100000

大致思路：

这个题刚开始想到判断A[i] + A[k] == 2*A[j]可以用FFT处理出任意两个数的和为T的有多少种, 然而发现这样难处理i, j, k的顺序

后来看了各路题解发现可以分块

考虑将整个数列分成K块, 每块 N / K 个连续的元素 (当然不一定都是N / K个, 数量级是这样)

那么考虑三个数的位置

如果三个数再同一块中, 直接对每一块进行枚举后两个数的位置然后查询当前块中需要成为等差数列需要的那个数有多少个, 时间复杂度是O(N/K * N/K *K)

如果三个数中有两个在同一块中, 枚举那两个数, 另外一个数可能在前面的块中也可能在后面的块中, 同上直接查询即可, 时间复杂度是 O(N/ K * N/K *K)

如果三个数再不同块中, 枚举中间那个数的在这当前块中的位置N/K种, 然后对于此块前后的部分作FFT得到前后各取一个数和是T有多少种, 对于N/K次直接查询即可, 记最大次数 为M = 1 << 16 (65536), 那么FFT一共进行K次, 时间复杂度是O(K*MlogM) 综上总体时间复杂度是O(N*N/K + K*MlogM), 去一个较为合适的K值即可

代码如下：

Result  :  Accepted     Memory  :  6451 KB     Time  :  27230 ms

/*
 * Author: Gatevin
 * Created Time:  2015/7/16 12:52:44
 * File Name: CodeChefCOUNTARI.cpp
 */
#include<iostream>
#include<sstream>
#include<fstream>
#include<vector>
#include<list>
#include<deque>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<bitset>
#include<algorithm>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cctype>
#include<cmath>
#include<ctime>
#include<iomanip>
using namespace std;
const double eps(1e-8);
typedef long long lint;

const double PI = acos(-1.0);

struct Complex
{
    double real, image;
    Complex(double _real, double _image)
    {
        real = _real;
        image = _image;
    }
    Complex(){}
};

Complex operator + (const Complex &c1, const Complex &c2)
{
    return Complex(c1.real + c2.real, c1.image + c2.image);
}

Complex operator - (const Complex &c1, const Complex &c2)
{
    return Complex(c1.real - c2.real, c1.image - c2.image);
}

Complex operator * (const Complex &c1, const Complex &c2)
{
    return Complex(c1.real *c2.real - c1.image*c2.image, c1.real*c2.image + c1.image*c2.real);
}

int rev(int id, int len)
{
    int ret = 0;
    for(int i = 0; (1 << i) < len; i++)
    {
        ret <<= 1;
        if(id & (1 << i)) ret |= 1;
    }
    return ret;
}

Complex A[1 << 16];
void FFT(Complex *a, int len, int DFT)
{
    for(int i = 0; i < len; i++)
        A[rev(i, len)] = a[i];
    for(int s = 1; (1 << s) <= len; s++)
    {
        int m = (1 << s);
        Complex wm = Complex(cos(DFT*2*PI/m), sin(DFT*2*PI/m));
        for(int k = 0; k < len; k += m)
        {
            Complex w = Complex(1, 0);
            for(int j = 0; j < (m >> 1); j++)
            {
                Complex t = w*A[k + j + (m >> 1)];
                Complex u = A[k + j];
                A[k + j] = u + t;
                A[k + j + (m >> 1)] = u - t;
                w = w*wm;
            }
        }
    }
    if(DFT == -1) for(int i = 0; i < len; i++) A[i].real /= len, A[i].image /= len;
    for(int i = 0; i < len; i++) a[i] = A[i];
    return;
}

int num[100010];
Complex L[1 << 16], R[1 << 16];
const int K = 44;
int before[30010], behind[30010], in[30010];//当前块前面的块, 后面的块, 当前块中的i的个数为before[i], behind[i], in[i]

int main()
{
    int N;
    while(~scanf("%d", &N))
    {
        int maxNum = 1;
        memset(before, 0, sizeof(before));
        memset(behind, 0, sizeof(behind));
        memset(in, 0, sizeof(in));
        for(int i = 1; i <= N; i++)
            scanf("%d", num + i), behind[num[i]]++, maxNum = max(maxNum, num[i]);
        int len = 1;
        while(len <= maxNum) len <<= 1;
        len <<= 1;//为FFT作准备, 确定多项式需要的最大次数
        
        int block = min(K, N);
        //int size = (N + block - 1) / block;//可以直接这么写的...
        int size = N / block;
        if(N % block != 0) size++;
        //这里之前写的block++逻辑错了, 不过block++取K = 100能AC..成功避开所有错误数据,什么人品..改了K值试了几下才发现
        
        lint ans = 0;
        for(int pos = 1; pos <= block; pos++)//当前第几块
        {
            int s = size*(pos - 1) + 1, e = size*pos;
            if(e > N) e = N;
            for(int t = s; t <= e; t++)//更新behind数组
                behind[num[t]]--;
            for(int i = s; i <= e; i++)
            {
                for(int j = i + 1; j <= e; j++)//枚举两个数
                {
                    int ak = 2*num[i] - num[j];
                    if(ak >= 1 && ak <= 30000)
                    {
                        ans += in[ak];//三个数再同一块中时, 枚举的是后两个数, 询问同一块的in数组
                        //枚举后两个数在同一块, 第一个数不在时
                        ans += before[ak];
                    }
                    ak = 2*num[j] - num[i];//枚举前两个数时, 第三个数再后面的块中
                    if(ak >= 1 && ak <= 30000)
                        ans += behind[ak];
                    
                }
                in[num[i]]++;
            }
            
            //接下来是三个数分布在三块中的情况, 枚举中间数, 对前后部分作FFT
            for(int i = 0; i <= maxNum; i++)
            {
                L[i] = Complex(before[i]*1., 0);
                R[i] = Complex(behind[i]*1., 0);
            }
            for(int i = maxNum + 1; i < len; i++)
                L[i] = R[i] = Complex(0, 0);
            FFT(L, len, 1); FFT(R, len, 1);
            for(int i = 0; i < len; i++)
                L[i] = L[i]*R[i];
            FFT(L, len, -1);
            for(int j = s; j <= e; j++)
                ans += (lint)(L[2*num[j]].real + 0.5);
            for(int t = s; t <= e; t++)//更新before数组, 还原in数组
                before[num[t]]++, in[num[t]]--;
        }
        printf("%lld\n", ans);
    }
    return 0;
}