Description

题库链接

给你一个长度为 \(n\) 的序列 \(A\) ，和一个数 \(x\) ，对于每个 \(i= 0\sim n\) ，求有多少个非空子区间满足恰好有 \(i\) 个数 \(<x\) 。

\(1 \leq n \leq 2 \cdot 10^5\)

Solution

我们用 \([A_i<x]\) 做一个前缀和，显然这是单调递增的。记前缀和为 \(i\) 的个数为 \(f_i\) 。

显然对于 \(i=k,k\neq 0\) 答案就是

\[\sum_{i=0}^{n-k}f(i+k)f(i)\]

我们令 \(g(i)=f(n-i)\) ，那么答案就是

\[\sum_{i=0}^{n}f(i+k)g(n-i)\]

卷一下就好了，不过注意到是非空子集，所以对于 \(k=0\) 时要特判。

Code

#include <bits/stdc++.h>

#define ll long long

#define dob complex<double>

using namespace std;

const double pi = acos(-1.0);

const int N = (200000<<2)+5;

int n, m, x, sum[N], u, R[N], L, len, cnt[N];

dob a[N], b[N];

void FFT(dob *A, int o) {

    for (int i = 0; i < len; i++) if (i < R[i]) swap(A[i], A[R[i]]);

    for (int i = 1; i < len; i <<= 1) {

        dob wn(cos(pi/i), sin(pi*o/i)), x, y;

        for (int j = 0; j < len; j += (i<<1)) {

            dob w(1, 0);

            for (int k = 0; k < i; k++, w *= wn) {

                x = A[j+k], y = w*A[j+k+i];

                A[j+k] = x+y, A[j+k+i] = x-y;

            }

        }

    }

}

void work() {

    scanf("%d%d", &n, &x); cnt[0]++;

    for (int i = 1; i <= n; i++) {

        scanf("%d", &u), sum[i] = sum[i-1]+(u < x);

        cnt[sum[i]]++;

    }

    for (int i = 0; i <= n; i++) a[i] = cnt[i];

    for (int i = 0; i <= n; i++) b[i] = cnt[n-i];

    m = (n<<1);

    for (len = 1; len <= m; len <<= 1) ++L;

    for (int i = 0; i < len; i++) R[i] = (R[i>>1]>>1|((i&1)<<(L-1)));

    FFT(a, 1), FFT(b, 1);

    for (int i = 0; i < len; i++) a[i] = a[i]*b[i];

    FFT(a, -1);

    printf("%I64d", (ll)(a[n].real()/len+0.5-n)>>1);

    for (int i = 1; i <= n; i++) printf(" %I64d" , (ll)(a[i+n].real()/len+0.5));

}

int main() {work(); return 0; }