bzoj3509: [CodeChef] COUNTARI 分块

bzoj3509: [CodeChef] COUNTARI

Description

给定一个长度为N的数组A[]，求有多少对i, j, k（1<=i<j<k<=N）满足A[k]-A[j]=A[j]-A[i]。

Input

第一行一个整数N（N<=10^5）。
接下来一行N个数A[i]（A[i]<=30000）。

Output

一行一个整数。

Sample Input

10
3 5 3 6 3 4 10 4 5 2

Sample Output

9

分析

这很FFT
考虑暴力枚举等比中项$j$。
一种做法是把序列两边权值生成函数之后卷积，用FFT优化之。
复杂度是$O(nmlogn)$显然过不了。
分块FFT，块内$O(B^2)$暴力即可。
暴力/FFT写得不好都会被卡。这里用共轭复数加速优化了一下。

代码

#include<bits/stdc++.h>
int ri() {
    char c = getchar(); int x = 0, f = 1; for(;c < '0' || c > '9'; c = getchar()) if(c == '-') f = -1;
    for(;c >= '0' && c <= '9'; c = getchar()) x = (x << 1) + (x << 3) - '0' + c; return x * f;
}
const int N = 1e5 + 10;
const double pi = acos(-1.0);
struct cp {
    double r, i;
    cp(double _r = 0, double _i = 0) : r(_r), i(_i) {}
    cp operator + (cp a) {return cp(r + a.r, i + a.i);}
    cp operator - (cp a) {return cp(r - a.r, i - a.i);}
    cp operator * (cp a) {return cp(r * a.r - i * a.i, r * a.i + i * a.r);}
}w[N], A[N], B[N];
cp conj(cp a) {return cp(a.r, -a.i);}
int R[N], l[N], r[N], a[N], Lc[N], Rc[N], b[N], n, L, mx; long long Ans;
void Pre(int m) {
    int x = 0; L = 1;
    for(;L < m; L <<= 1, ++x) ;
    for(int i = 1;i < L; ++i) R[i] = (R[i >> 1] >> 1) | ((i & 1) << x - 1);
    for(int i = 0;i < L; ++i) w[i] = cp(cos(2 * pi * i / L), sin(2 * pi * i / L));
}
void FFT(cp *F) {
    for(int i = 0;i < L; ++i) if(i < R[i]) std::swap(F[i], F[R[i]]);
    for(int i = 1, d = L >> 1;i < L; i <<= 1, d >>= 1)
        for(int j = 0;j < L; j += (i << 1)) {
            cp *l = F + j, *r = F + j + i, *p = w, tp;
            for(int k = 0;k < i; ++k, ++l, ++r, p += d)
                tp = *r * *p, *r = *l - tp, *l = *l + tp;
        }
}
void DFT() {
	static cp C[N];
	for(int i = 0;i < L; ++i) C[i] = cp(Lc[i], Rc[i]);
	FFT(C);
	for(int i = 0;i < L; ++i) {
		int j = (L - i) & L - 1;
		A[i] = (C[i] + conj(C[j])) * cp(0.5, 0.0);
		B[i] = (C[i] - conj(C[j])) * cp(0.0, -0.5);
	}
}
void Work(int x) {
	mx = -1;
    for(int i = 1;i < l[x]; ++i) ++Lc[a[i]], mx = std::max(mx, a[i]);
    for(int i = r[x] + 1;i <= n; ++i) ++Rc[a[i]], mx = std::max(mx, a[i]);
    Pre(mx << 1); DFT();
    for(int i = 0;i < L; ++i) A[i] = A[i] * B[i];
    FFT(A); for(int i = 1;i < L >> 1; ++i) std::swap(A[i], A[L - i]);
    for(int i = l[x]; i <= r[x]; ++i) Ans += (long long)(A[a[i] << 1].r / L + 0.5);
    for(int i = 0;i <= mx; ++i) Lc[i] = Rc[i] = 0;
}
void Force() {
	for(int i = 1;i <= n; ++i) ++Rc[a[i]];
	for(int x = 1;x <= b[n]; ++x) {
		for(int i = l[x];i <= r[x]; ++i) --Rc[a[i]];
		for(int i = l[x];i <= r[x]; ++i) {
			for(int j = i + 1, v;j <= r[x]; ++j) {
				if((v = (a[i] << 1) - a[j]) >= 0) Ans += Lc[v];
				if((v = (a[j] << 1) - a[i]) >= 0) Ans += Rc[v];
			}
			++Lc[a[i]];
		}
	}
	for(int i = 0;i <= mx; ++i) Lc[i] = Rc[i] = 0;
}
int main() {
    n = ri(); int B = 2500;
    for(int i = 1;i <= n; ++i) mx = std::max(mx, a[i] = ri());
    for(int i = 1;i <= n; ++i) {
        b[i] = (i - 1) / B + 1; if(!l[b[i]]) l[b[i]] = i; r[b[i]] = i;
    }
    Force();
    for(int i = 2;i < b[n]; ++i) Work(i);
    printf("%lld\n", Ans);
    return 0;
}