백준 9426 - 중앙값 측정

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

 

지문의 힌트에서 (K + 1) / 2번째로 작은 숫자를 찾으면 된다고 알려주고 있다.

 

세그먼트 트리를 통해 (K + 1) / 2번째로 작은 숫자를 찾는 연산을 구현하면 간단히 해결할 수 있다.

 

중앙값을 구하는 구간의 길이는 K라고 명시하고 있으므로 슬라이딩 윈도우를 통해 구간을 벗어나는 원소나 구간에 속한 원소만을 update 연산으로 세그먼트 트리에 반영해주면 된다.

 

전체 코드

#include <bits/stdc++.h>

using namespace std;

int n, k;

inline int mid(int s, int e) { return s + (e - s) / 2; }

void update(vector<int> &tree, int node, int start, int end, int idx, int diff)
{
    if (idx < start || idx > end) return;

    tree[node] += diff;

    if (start != end)
    {
        int m = mid(start, end);
        update(tree, node * 2, start, m, idx, diff);
        update(tree, node * 2 + 1, m + 1, end, idx, diff);
    }
}

int query(vector<int> &tree, int node, int start, int end, int kth)
{
    if (start == end) return end;
    int m = mid(start, end);
    int left_val = tree[node * 2];

    if (kth <= left_val) return query(tree, node * 2, start, m, kth);
    return query(tree, node * 2 + 1, m + 1, end, kth - left_val);
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);

    cin >> n >> k;

    vector<int> segtree(65535 * 4);
    deque<int> dq;
    long long res = 0;

    for (int i = 0; i < n; ++i)
    {
        int x;
        cin >> x;

        dq.push_back(x);
        update(segtree, 1, 0, 65535, x, 1);

        if (dq.size() == k)
        {
            res += query(segtree, 1, 0, 65535, (k + 1) / 2);
            int target = dq.front();
            dq.pop_front(); 
            update(segtree, 1, 0, 65535, target, -1);
        }
            
    }

    cout << res;
}