Codeforces 993E Nikita and Order Statistics (FFT)

Codeforces Round #488 by NEAR (Div. 1)

题目链接：

Nikita and Order Statistics

CF993E Nikita and Order Statistics

Input

Output

Examples

output

output

output

Solution

题意

给定长度为 \(n\) 的数组，对于 \(k \in [0, n]\)，求有多少个区间满足区间内恰好有 \(k\) 个数比 \(x\) 小。

题解

FFT

首先，对于数组中的每一个数，如果比 \(x\) 小，那么把这个数置为 \(1\)，否则置为 \(0\)。然后求该数组的前缀和 \(s\)。

设 \(f[i]\) 表示前缀和等于 \(i\) 的个数。那么 \(ans[k] = \sum_{i=k}^nf[i]f[i-k] = \sum_{i=k}^nf[i]f[n-(n+k-i)] = \sum_{i+j=n+k}f[i]f[n-j]\)。

令 \(g[i] = f[n - i]\)，则 \(ans[k] = \sum_{i+j=n+k}f[i]g[j]\)。

问题就转化为求 \(f\) 与 \(g\) 的卷积。

注意 \(k = 0\) 的情况要特判。此时会有 \(n + 1\) 次自己和自己算到一起，也就是空串的情况。所以要对 \(ans[0]\) 减去 \(n + 1\)。减完后还要除以 \(2\)，因为 \(k \neq 0\) 时，由于 \(s\) 单调不减，因此 \(s[i] \ge s[i - k]\)。此时一定保证区间的左端点在前，右端点在后。而 \(k = 0\) 时，不能保证左端点一定在前，右端点一定在后，所以要除以 \(2\)。

Code

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const double PI = acos(-1);
const int maxn = 1e6 + 10;
typedef complex<double> Complex;

ll n, x;
ll s[maxn], cnt[maxn];
Complex f[maxn], g[maxn];
ll ans[maxn];
int bit = 2, rev[maxn];

void get_rev(){
    memset(rev, 0, sizeof(rev));
    while(bit <= n * 2) bit <<= 1;
    for(int i = 0; i < bit; ++i) {
        rev[i] = (rev[i >> 1] >> 1) | (bit >> 1) * (i & 1);
    }
}

void FFT(Complex *a, int op) {
    for(int i = 0; i < bit; ++i) {
        if(i < rev[i]) swap(a[i], a[rev[i]]);
    }
    for(int mid = 1; mid < bit; mid <<= 1) {
        Complex wn = Complex(cos(PI / mid), op * sin(PI / mid));
        for(int j = 0; j < bit; j += mid<<1) {
            Complex w(1, 0);
            for(int k = 0; k < mid; ++k, w = w * wn) {
                Complex x = a[j + k], y = w * a[j + k + mid];
                a[j + k] = x + y, a[j + k + mid] = x - y;
            }
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> x;
    ++cnt[0];
    for(int i = 1; i <= n; ++i) {
        int a;
        cin >> a;
        s[i] = s[i - 1] + (a < x);
        ++cnt[s[i]];
    }
    for(int i = 0; i <= n; ++i) {
        f[i] = cnt[i];
        g[n - i] = cnt[i];
    }
    get_rev();
    FFT(f, 1), FFT(g, 1);
    for(int i = 0; i < bit; ++i) {
        f[i] = f[i] * g[i];
    }
    FFT(f, -1);
    for(int i = 0; i <= n; ++i) {
        ans[i] = (ll)(f[n + i].real() / bit + 0.5);
    }
    cout << (ans[0] - n - 1) / 2 << " ";
    for(int i = 1; i <= n; ++i) {
        cout << ans[i] << " ";
    }
    cout << endl;
    return 0;
}