1572번: 중앙값

https://www.acmicpc.net/problem/1572


$O(n\lg n)$

#include<cstdio>
#include<set>
using namespace std;
int n, k, a[250000];
long long r;
struct st {
    multiset<int> ms;
    multiset<int>::iterator it, t;
    int lc, rc;
    void kb() {
        if (lc > rc) it--, lc--, rc++;
        if (lc + 1 < rc) it++, lc++, rc--;
    }
    void insert(int x) {
        ms.insert(x);
        if (ms.size() == 1) it = ms.begin();
        else x < *it ? lc++ : rc++;
        kb();
    }
    void erase(int x) {
        t = ms.lower_bound(x);
        if (t == it) it++, rc -= ms.size()>1;
        else x > *it ? rc-- : lc--;
        ms.erase(t);
        kb();
    }
    int get() { return *it; }
}mid;
int main() {
    scanf("%d%d", &n, &k);
    for (int i = 0; i < n; i++) scanf("%d", a + i);
    for (int i = 0; i < k - 1; i++) mid.insert(a[i]);
    for (int i = k - 1; i < n; i++) {
        mid.insert(a[i]);
        r += mid.get();
        mid.erase(a[i - k + 1]);
    }
    printf("%lld", r);
    return 0;
}