[백준 9426] 중앙값 측정

문제

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

 

알고리즘

세그먼트 트리

 

풀이

$N$초동안 온도를 기록하는데 최근 $K$초의 중앙값들을 모두 더하는 것을 요구하는 문제입니다.

 

이 문제의 경우는 이전의 값을 삭제하는 연산 없이 계속 추가한 후 중앙값만 처리하면 되므로 최대 힙, 최소 힙 두 개로 $O(NlogN)$에 해결 가능했으나, 가장 최근$K$초의 중앙값만을 계속 구해야 하므로 최근 K초를 제외한 온도를 삭제할 수 있는 자료구조인 세그먼트 트리를 활용하겠습니다.

 

문제에서 요구하는 핵심쿼리는 중앙값 처리입니다. 트리에서의 인덱스는 온도를 의미합니다. 온도의 개수를 누적시킨다면, 인덱스에 따라 값들이 오름차순으로 정렬되어 있으므로, 이분 탐색을 통해 중앙값을 빠르게 처리할 수 있습니다.

즉 세그먼트 트리중에서도 펜윅트리를 사용합시다.

 

펜윅트리를 사용할 때 주의할 것은 인덱스를 사용하는 연산이므로 0인덱스는 사용하지 못합니다. 문제에서 주어지는 온도의 범위는 0~65535이므로 1씩 더해 1~65536으로 우선 계산한 후 이분 탐색을 통해 중앙값을 구했다면 다시 1을 감소시켜주도록 합시다. 시간복잡도는 $O(NlogN)$입니다.

 

코드

#include <bits/stdc++.h>
#define rep(i,n) for(int i=0;i<n;++i)
#define REP(i,n) for(int i=1;i<=n;++i)
#define FAST cin.tie(NULL);cout.tie(NULL); ios::sync_with_stdio(false)
using namespace std;

const int MX = 65535 + 1;

int N, K;
long long ans;
int fenwick[MX], arr[250001];

void update(int idx,int val) {
    while (idx <= MX) {
        fenwick[idx] += val;
        idx += idx & -idx;
    }
}

int query(int idx) {
    int ret = 0;
    while (idx) {
        ret += fenwick[idx];
        idx -= idx & -idx;
    }
    return ret;
}

int main() {
    FAST;
    cin >> N >> K;

    REP(i, N) {
        cin >> arr[i];
        update(++arr[i], 1);
        if (i >= K) {
            int lo = 1;
            int hi = MX;
            int best;
            while (lo <= hi) {
                int m = (lo + hi) / 2;
                if (query(m) >= (K + 1) / 2) {
                    best = m;
                    hi = m - 1;
                }
                else lo = m + 1;
            }
            ans += best - 1;
            update(arr[i - K + 1], -1);
        }
    }
    cout << ans;
    return 0;
}