CodeChef - COUNTARI Arithmetic Progressions (FFT)

Posted 2021-01-09 xiuwenli

题意：求一个序列中，有多少三元组$（i,j,k）i<j<k $ 满足(A_i + A_k = 2*A_i) 构成等差数列。

https://www.cnblogs.com/xiuwenli/p/9719425.html

在该题的基础上加了i<j<k的限制。不能单纯地FFT枚举结果。
考虑分块搭配FFT结果问题。
将原序列分成若干块，设每个块的大小为(B)。则构成等差数列的三元组有三种构成情况。
1）三个数都在一个块内。这种情况对每个块独立处理。二重循环枚举(A_i 、A_k)，不断更新(i,k)之间满足条件的(A_j)；
2）两个数在一个块内。那么有两种可能：(A_i、A_j)在一个块内，或(A_j、A_k)在一个块内。此时需要之前所有块的信息和后面所有块的信息。
3）三个数在不同块中。此时情况和不带限制的相似。统计出之前所有块的信息和之后所有块的信息，FFT计算出和的系数，扫一遍块内数据统计结果。

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int MAXN = 1e5 + 10;
const double PI = acos(-1.0);
struct Complex{
    double x, y;
    inline Complex operator+(const Complex b) const {
        return (Complex){x +b.x,y + b.y};
    }
    inline Complex operator-(const Complex b) const {
        return (Complex){x -b.x,y - b.y};
    }
    inline Complex operator*(const Complex b) const {
        return (Complex){x *b.x -y * b.y,x * b.y + y * b.x};
    }
} va[MAXN * 2 + MAXN / 2], vb[MAXN * 2 + MAXN / 2];
int lenth = 1, rev[MAXN * 2 + MAXN / 2];
int N, M;   // f 和 g 的数量
    //f g和 的系数
    // 卷积结果
    // 大数乘积
int f[MAXN],g[MAXN];
vector<LL> conv;
vector<LL> multi;
void debug(){for(int i=0;i<conv.size();++i) cout<<conv[i]<<" ";cout<<endl;}
//f g
void init()
{
    int tim = 0;
    lenth = 1;
    conv.clear(), multi.clear();
    memset(va, 0, sizeof va);
    memset(vb, 0, sizeof vb);
    while (lenth <= N + M - 2)
        lenth <<= 1, tim++;
    for (int i = 0; i < lenth; i++)
        rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (tim - 1));
}
void FFT(Complex *A, const int fla)
{
    for (int i = 0; i < lenth; i++){
        if (i < rev[i]){
            swap(A[i], A[rev[i]]);
        }
    }
    for (int i = 1; i < lenth; i <<= 1){
        const Complex w = (Complex){cos(PI / i), fla * sin(PI / i)};
        for (int j = 0; j < lenth; j += (i << 1)){
            Complex K = (Complex){1, 0};
            for (int k = 0; k < i; k++, K = K * w){
                const Complex x = A[j + k], y = K * A[j + k + i];
                A[j + k] = x + y;
                A[j + k + i] = x - y;
            }
        }
    }
}
void getConv(){             //求多项式
    init();
    for (int i = 0; i < N; i++)
        va[i].x = f[i];
    for (int i = 0; i < M; i++)
        vb[i].x = g[i];
    FFT(va, 1), FFT(vb, 1);
    for (int i = 0; i < lenth; i++)
        va[i] = va[i] * vb[i];
    FFT(va, -1);
    for (int i = 0; i <= N + M - 2; i++)
        conv.push_back((LL)(va[i].x / lenth + 0.5));
}

int a[MAXN];
int cnt[30005];
int cnt2[30005];

int main()
{
    #ifndef ONLINE_JUDGE
        freopen("in.txt","r",stdin);
        freopen("out.txt","w",stdout);
    #endif
    int n;
    while(scanf("%d",&n)==1){
        memset(cnt,0,sizeof(cnt));
        for(int i=0;i<n;++i){
            scanf("%d",&a[i]);
        }   
        int block=(int)(sqrt(n)+1e-6)*5;
        LL res=0;
        //1) 三个数都在一个块内
        for(int b=0;b<n;b+=block){
            for(int i=0; i<block && b+i<n;++i){
                for(int j=i+1;j<block && b+j<n;++j){
                    int sum = a[b+i] + a[b+j];
                    if((sum&1)==0){
                        res += cnt[sum>>1];        
                    }
                    ++cnt[a[b+j]];
                }
                for(int j=i+1;j<block && b+j<n ;++j)
                    --cnt[a[b+j]];
            }
        }

        //2) 两个数在一个块内
        memset(cnt,0,sizeof(cnt));          //后面所有块的信息
        memset(cnt2,0,sizeof(cnt2));        //前面所有块的信息
        for(int i=0;i<n;++i){
            cnt[a[i]]++;
        }
        for(int b = 0;b<n;b+=block){
            for(int i=0;i<block && b+i<n;++i){          //将该块的数据抹去
                cnt[a[b+i]]--;
            }    
            for(int i=0;i<block && b+i<n;++i){
                for(int j=i+1;j<block && b+j<n;++j){
                    int cha = 2*a[b+i] - a[b+j];
                    if(cha>0 && cha<=30000) res += cnt2[cha];
                    int sum = 2*a[b+j] - a[b+i];
                    if(sum>0 && sum<=30000) res += cnt[sum];
                }
            }
            for(int i=0;i<block && b+i<n;++i){          //将该块的数据更新到 cnt2
                ++cnt2[a[b+i]];
            }
        }
        // 3) 三个数在不同块中 
        memset(cnt,0,sizeof(cnt));          //后面所有块的信息
        memset(cnt2,0,sizeof(cnt2));        //前面所有块的信息
        for(int i=0;i<n;++i){
            cnt[a[i]]++;
        }
        for(int b = 0;b<n;b+=block){
            for(int i=0;i<block && b+i<n;++i){          //将该块的数据抹去
                cnt[a[b+i]]--;
            }   
            if(b>1){
                N =M = 30001; 
                for(int i=0;i<=30000;++i){
                    f[i] = cnt[i];
                    g[i] = cnt2[i];
                }
                getConv();
                int sz = conv.size();
                for(int i=0;i<block &&b+i<n;++i){
                    if(a[b+i]*2>=sz) continue;
                    res += conv[a[b+i]*2];
                }
            }
            for(int i=0;i<block && b+i<n;++i){          //将该块的数据更新到 cnt2
                ++cnt2[a[b+i]];
            }
        }
        printf("%lld
",res);
    }
    return 0;
}