중앙값

백준 1572번: 중앙값

아이디어

사탕상자 문제랑 비슷하다. 구간에서 순위를 구하는 문제는 나올 수 있는 전체 숫자 사이즈만큼 범위로 잡아서 트리를 만들고, 숫자가 나올 때마다 개수를 늘려주고, 해당 구간에서 개수의 합을 반환하는 쿼리를 짜면 된다.
최근 숫자 K개중에 중앙값을 구해야 하기 때문에 큐에서 숫자 K개를 유지하며 push, pop 한다.

코드

#include <bits/stdc++.h>

using namespace std;

const int MAX = 65536;
int N, K;

vector<int> tree(4*MAX, 0);

void update(int start, int end, int idx, int diff, int node) {
    if (idx < start || end < idx) return;
    tree[node] += diff;
    if (start == end) return;
    update(start, (start+end)/2, idx, diff, node*2);
    update((start+end)/2+1, end, idx, diff, node*2+1);
}

int getmid(int start, int end, int midx, int node) {
    if (start == end) {
        return start;
    }
    if (tree[node*2] < midx) {
        return getmid((start+end)/2+1, end, midx-tree[node*2], node*2+1);
    }
    else {
        return getmid(start, (start+end)/2, midx, node*2);
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> N >> K;
    int midx = K/2+1;
    if (K%2 == 0) midx--;
    int x;
    long long ans = 0;
    queue<int> q;
    for (int i = 0; i < K; i++) {
        cin >> x;
        q.push(x);
        update(0, MAX-1, x, 1, 1);
    }
    ans += getmid(0, MAX-1, midx, 1);
    for (int i = K; i < N; i++) {
        update(0, MAX-1, q.front(), -1, 1);
        q.pop();
        cin >> x;
        q.push(x);
        update(0, MAX-1, x, 1, 1);
        ans += getmid(0, MAX-1, midx, 1);
    }
    cout << ans;
    return 0;
}

여담

비슷한거 풀었는데 왜 오래걸렸지? 입력 값 범위만큼 트리를 만드는걸 항상 고려해봐야겠다.