[CodeForces] CF993E Nikita and Order Statistics

Posted 2021-03-13 hock

题目大意

给定长度为(n)的序列(a_i)和(x)
你需要求出(ans_i)
(ans_i)表示有多少个区间满足有(i)个数小于(x)

题解

很明显 这个数列只有两类数 大于等于(x)和小于(x)的两种数 那么不妨(geq x)的设为(0) 其他设为(1) 那么(ans_i)表示的就是有多少个区间满足(s_k - s_j = i) (s)是前缀和这里不再赘述

(O(n^2))暴力

(s_i - s_j = k) 这个很明显可以用(two pointer)来做(O(n))扫描整个数组但是每个都要计算的话 那么时间复杂度为(O(n^2))

(O(n log n))正解

(s_i - s_j = k) 我们移项(s_i - k = s_j)我们对于每个(i)我们只要求多少个(s_i - k)即可 所以(ans_i = sum_{k}^n f(i)*f(i-k))
令(g(i) = f(i - k)) 那么(ans_i = sum f(i) * g(i))
我们可以用(FFT)来优化这个算法
我们可以令(g(i) = f(n-i)) 那么多项式乘法后(ans_i)即等于(h(n-i))

Code

#include <cstdio>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <set>
#include <map>
#include <queue>

using namespace std;

template <typename T> void chkmax (T &x, T y) {x = x > y ? x : y;}
template <typename T> void chkmin (T &x, T y) {x = x > y ? y : x;}

typedef long long ll;

const int INF = 2139062143;

#define DEBUG(x) std::cerr << #x << " = " << x << std::endl

template <typename T> void read (T &x) {
    x = 0; bool f = 1; char ch;
    do {ch = getchar(); if (ch == '-') f = 0;} while (ch > '9' || ch < '0');
    do {x = x * 10 + ch - '0'; ch = getchar();} while (ch >= '0' && ch <= '9');
    x = f ? x : -x;
}

template <typename T> void write (T x) {
    if (x < 0) x = ~x + 1, putchar ('-');
    if (x > 9) write (x / 10);
    putchar (x % 10 + '0');
}

const int N = 2e5 + 7;
const double PI = acos (-1.0);

struct complex {
    double x, y;
    complex (double xx = 0, double yy = 0) {x = xx, y = yy;}
} f[N], g[N];

complex operator + (complex a, complex b) {return complex (a.x + b.x, a.y + b.y);}
complex operator - (complex a, complex b) {return complex (a.x - b.x, a.y - b.y);}
complex operator * (complex a, complex b) {return complex (a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);}

int n, x, l, r[N], s[N];

void FFT (complex *a, int limit, int opt) {
    for (int i = 0; i < limit; i ++ ) if (i < r[i]) swap (a[i], a[r[i]]);
    for (int i = 1; i < limit; i <<= 1) {
        complex Wn = complex (cos (PI / i), opt * sin (PI / i));
        for (int j = 0; j < limit; j += (i << 1)) {
            complex w = complex (1, 0);
            for (int k = 0; k < i; k ++, w = w * Wn) {
                complex p = a[j + k], q = w * a[j + k + i];
                a[j + k] = p + q; a[j + k + i] = p - q;
            }
        }
    }
    if (opt == -1) for (int i = 0; i < limit; i ++ ) a[i].x = a[i].x / limit;
}

inline ll calc() {
    ll res = 0;
    for (int i = 1, p = 0; i <= n; ++ i ) {
        if (s[i] != s[i - 1]) p = i;
        res += i - p;
    }
    return res;
}
int main () {
    read (n); read (x);
    f[0].x ++ ;
    for (int i = 1, k; i <= n; i ++ ) read (k), s[i] = s[i - 1] + (k < x), f[s[i]].x ++ ;
    for (int i = 0; i <= n; i ++ ) g[i] = f[n - i];
    int limit = 1; while (limit <= n + n) limit <<= 1, l ++ ;
    for (int i = 0; i < limit; i ++ ) r[i] = (r[i >> 1] >> 1) | ((i & 1) << (l - 1));

    FFT (f, limit, 1); FFT (g, limit, 1);
    for (int i = 0; i < limit; i ++ ) f[i] = f[i] * g[i];
    FFT (f, limit, -1);
    write (calc ()); putchar (' ');
    for (int i = 1; i <= n; i ++ ) printf ("%lld ", (ll) (f[n - i].x + 0.5));
    return 0;
}