[Codeforces 993E]Nikita and Order Statistics

Description

题库链接

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

(1 leq n leq 2 cdot 10^5)

Solution

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

显然对于 (i=k,k eq 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; }