bzoj 3509: [CodeChef] COUNTARI] [分块 生成函数]

3509: [CodeChef] COUNTARI

题意：统计满足\(i<j<k, 2*a[j] = a[i] + a[k]\)的个数

---

\(2*a[j]\)不太好处理，暴力fft不如直接暴力

考虑分块

每个块用生成函数统计j在块中ik在两边的块中的

有两个在块中或者三个都在暴力统计，实时维护两边权值出现次数

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <ctime>
using namespace std;
typedef long long ll;
const int N = (1<<16) + 5, M = 1e5+5;
const double PI = acos(-1.0);
inline int read() {
    char c=getchar(); int x=0,f=1;
    while(c<'0'||c>'9') {if(c=='-')f=-1;c=getchar();}
    while(c>='0'&&c<='9') {x=x*10+c-'0';c=getchar();}
    return x*f;
}

struct meow{
    double x, y;
    meow(double a=0, double b=0):x(a), y(b){}
};
meow operator +(meow a, meow b) {return meow(a.x+b.x, a.y+b.y);}
meow operator -(meow a, meow b) {return meow(a.x-b.x, a.y-b.y);}
meow operator *(meow a, meow b) {return meow(a.x*b.x-a.y*b.y, a.x*b.y+a.y*b.x);}
meow conj(meow a) {return meow(a.x, -a.y);}
typedef meow cd;

namespace fft {
    int n, rev[N];
    cd omega[N], omegaInv[N];
    void init(int lim) {
        n = 1; while(n < lim) n <<= 1;
        for(int i=0; i<n; i++) {
            omega[i] = cd(cos(2*PI/n*i), sin(2*PI/n*i));
            omegaInv[i] = conj(omega[i]);
        }
    }
    void dft(cd *a, int n, int flag) {
        cd *w = flag == 1 ? omega : omegaInv;
        int k = 0; while((1<<k) < n) k++;
        for(int i=0; i<n; i++) {
            rev[i] = (rev[i>>1]>>1) | ((i&1)<<(k-1));
            if(i < rev[i]) swap(a[i], a[rev[i]]);
        }
        for(int l=2; l<=n; l<<=1) {
            int m = l>>1;
            for(cd *p = a; p != a+n; p += l) 
                for(int k=0; k<m; k++) {
                    cd t = w[n/l*k] * p[k+m];
                    p[k+m] = p[k] - t;
                    p[k] = p[k] + t;
                }
        }
        if(flag == -1) for(int i=0; i<n; i++) a[i].x /= n;
    }
    void mul(cd *a, cd *b) {
        dft(a, n, 1); dft(b, n, 1);
        for(int i=0; i<n; i++) a[i] = a[i] * b[i];
        dft(a, n, -1);
    }
}

int n, a[M], c1[N], c2[N], block, m;
cd p[N], q[N];
ll ans;
int main() {
    freopen("in", "r", stdin);
    n=read();
    block = min(n, 8 * (int) sqrt(n));
    for(int i=1; i<=n; i++) a[i] = read(), c2[a[i]]++, m = max(m, a[i]);
    fft::init(m+m+1);
    for(int l=1; l<=n; l+=block) {
        int r = min(n, l+block-1);
        for(int i=l; i<=r; i++) c2[a[i]]--;
        for(int i=l; i<=r; i++) {
            for(int j=i+1; j<=r; j++) {
                int t = (a[j]<<1) - a[i];
                if(t >= 0) ans += c2[t];
                t = (a[i]<<1) - a[j];
                if(t >= 0) ans += c1[t];
            }
            c1[a[i]]++;
        }
        //printf("hi [%d, %d] %lld\n", l, r, ans);
    } 
    for(int l=1; l<=n; l+=block) { //printf("l %d\n", l);
        int r = min(n, l+block-1);
        memset(p, 0, sizeof(p)); memset(q, 0, sizeof(q));
        for(int i=1; i<l; i++) p[a[i]].x ++;
        for(int i=r+1; i<=n; i++) q[a[i]].x ++;
        fft::mul(p, q);
        for(int i=l; i<=r; i++) ans += (ll) floor(p[a[i]<<1].x + 0.5);
    }
    printf("%lld", ans);
    //printf("\n%lf", (double) clock() / CLOCKS_PER_SEC);
}