題目鏈接

BZOJ3509

題解

化一下式子，就是
\[2A[j] = A[i] + A[k]\]
所以我們對一個位置兩邊的數構成的生成函數相乘即可
但是由於這樣做是\(O(n^2logn)\)的，我們考慮如何優化

顯然可以分塊做，我們不對所有數左右求卷積，只對\(B\)個塊左右做，這樣\(i\)和\(k\)都在塊外的情況就可以統計出來
\(i\)或\(k\)在塊內的情況可以暴力掃一遍計算
復雜度\(O(Bnlogn + nB)\)
經計算\(B = \sqrt{nlogn}\)最優
但考慮到\(fft\)的常數問題，\(B = 2000\)左右比較合理
復雜度就大概是\(O((nlogn)^{\frac{3}{2}})\)

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#define REP(i,n) for (register int i = 1; i <= (n); i++)
#define mp(a,b) make_pair<int,int>(a,b)
#define cp pair<int,int>
#define LL long long int
 

#define res register
using namespace std;
const int maxn = 400005,maxm = 4005,INF = 1000000000;
inline int read(){
    int out = 0,flag = 1; char c = getchar();
    while (c < 48 || c > 57){if (c == ‘-‘) flag = -1; c = getchar();}
    while (c >= 48 && c <= 57){out = (out << 3) + (out << 1
 
) + c - 48; c = getchar();}
    return out * flag;
}
struct E{
    double a,b;
    E(){}
    E(double x,double y):a(x),b(y) {}
    E(int x,int y):a(x),b(y) {}
    inline E operator =(const int& b){
        this->a = b; this->b = 0;
        return *this;
    }
    inline E operator =(const double& b){
        this->a = b; this->b = 0;
        return *this;
    }
    inline E operator /=(const double& b){
        this->a /= b; this->b /= b;
        return *this;
    }
};
inline E operator *(const E& a,const E& b){
    return E(a.a * b.a - a.b * b.b,a.a * b.b + a.b * b.a);
}
inline E operator *=(E& a,const E& b){
    return a = E(a.a * b.a - a.b * b.b,a.a * b.b + a.b * b.a);
}
inline E operator +(const E& a,const E& b){
    return E(a.a + b.a,a.b + b.b);
}
inline E operator -(const E& a,const E& b){
    return E(a.a - b.a,a.b - b.b);
}
const double pi = acos(-1);
int R[maxn];
void fft(E* a,int n,int f){
    for (res int i = 0; i < n; i++) if (i < R[i]) swap(a[i],a[R[i]]);
    for (res int i = 1; i < n; i <<= 1){
        E wn(cos(pi / i),f * sin(pi / i));
        for (res int j = 0; j < n; j += (i << 1)){
            E w(1,0),x,y;
            for (res int k = 0; k < i; k++,w = w * wn){
                x = a[j + k],y = w * a[j + k + i];
                a[j + k] = x + y; a[j + k + i] = x - y;
            }
        }
    }
    if (f == -1) for (int i = 0; i < n; i++) a[i] /= n;
}
E A[maxn],C[maxn];
int N,B,val[maxn],b[maxn],bl[maxm],br[maxm],bi;
LL ans;
void work1(){
    for (res int x = 2; x < bi; x++){
        int deg1 = 0,deg2 = 0;
        for (res int i = 1; b[i] != x; i++)
            A[val[i]].a++,deg1 = max(deg1,val[i]);
        for (res int i = N; b[i] != x; i--)
            C[val[i]].a++,deg2 = max(deg2,val[i]);
        int n = 1,L = 0;
        while (n <= deg1 + deg2) n <<= 1,L++;
        for (res int i = 1; i < n; i++) R[i] = (R[i >> 1] >> 1) | ((i & 1) << (L - 1));
        fft(A,n,1); fft(C,n,1);
        for (res int i = 0; i < n; i++) A[i] *= C[i];
        fft(A,n,-1);
        for (res int i = bl[x]; i <= br[x]; i++)
            ans += (LL)floor(A[val[i] << 1].a + 0.1);
        for (res int i = 0; i < n; i++) A[i] = C[i] = 0.0;
    }
}
int bac[maxn],M;
void work2(){
    for (res int i = 1; i <= N; i++){
        for (res int j = i + 1; b[j] == b[i]; j++){
            if (val[j] >= val[i] && val[i] >= val[j] - val[i])
                ans += bac[val[i] - (val[j] - val[i])];
            if (val[j] < val[i] && val[i] + val[i] - val[j] <= 30000)
                ans += bac[val[i] + val[i] - val[j]];
        }
        bac[val[i]]++;
    }
    REP(i,N) bac[val[i]]--;
    for (res int i = N; i; i--){
        for (res int j = i - 1; b[j] == b[i]; j--){
            if (val[j] >= val[i] && val[i] >= val[j] - val[i])
                ans += bac[val[i] - (val[j] - val[i])];
            if (val[j] < val[i] && val[i] + val[i] - val[j] <= 30000)
                ans += bac[val[i] + val[i] - val[j]];
        }
        bac[val[i]]++;
    }
    REP(i,N) bac[val[i]]--;
    for (res int x = 1; x <= bi; x++){
        for (res int i = bl[x]; i <= br[x]; i++){
            for (res int j = i + 1; j <= br[x]; j++){
                if (val[j] >= val[i] && val[i] >= val[j] - val[i])
                    ans -= bac[val[i] - (val[j] - val[i])];
                if (val[j] < val[i] && val[i] + val[i] - val[j] <= 30000)
                    ans -= bac[val[i] + val[i] - val[j]];
            }
            bac[val[i]]++;
        }
        for (int i = bl[x]; i <= br[x]; i++) bac[val[i]]--;
    }
}
int main(){
    N = read(); B = 2000;
    REP(i,N){
        val[i] = read(),b[i] = i / B + 1,M = max(M,val[i]);
        if (b[i] != b[i - 1]) br[b[i - 1]] = i - 1,bl[b[i]] = i;
    }
    br[b[N]] = N; bi = b[N];
    work1();
    work2();
    printf("%lld\n",ans);
    return 0;
}

BZOJ3509 [CodeChef] COUNTARI 【分塊 + fft】