백준 1655번 가운데를 말해요

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

1655번: 가운데를 말해요

문제

N개의 정수를 입력받아 1번째, 2번째, ... N번째 정수까지만 보았을 때 각각의 중앙값들을 출력하는 문제입니다. 

풀이

단순히 배열을 사용하여 입력받을 때마다 배열을 정렬한다면 O(N^2)이므로 시간초과일 것입니다.

힙을 이용하여 O(NlogN)으로 해결할 수 있습니다.

지금까지 입력받은 수 중 절반의 작은 수들은 최대 힙에, 나머지 절반의 큰 수들은 최소 힙에 저장합니다. 

이렇게 한다면 중간값은 두 힙의 top값 중 하나이므로 적절한 비교를 통해 중간값을 찾을 수 있습니다.

정수가 N번 입력될 때마다 힙에서 O(logN)인 삽입/삭제가 일어나므로 총 O(NlogN)의 시간복잡도를 가집니다. 

코드

#include <cstdio>
#include <queue>
using namespace std;
 
priority_queue<int> small, large;
 
int main() {
    int n, got, sma_size = 0, lar_size = 0;
    scanf("%d", &n);
    for (int i = 0; i < n; i++) {
        scanf("%d", &got);
        if (i) {
            if (got > small.top()) {
                large.push(-got);
                lar_size++;
                if (lar_size > sma_size + 1) {
                    small.push(-large.top());
                    large.pop();
                    sma_size++;
                    lar_size--;
                }
            }
            else {
                small.push(got);
                sma_size++;
                if (sma_size > lar_size + 1) {
                    large.push(-small.top());
                    small.pop();
                    lar_size++;
                    sma_size--;
                }
            }
        }
        else {
            small.push(got);
            sma_size++;
        }
        if (i % 2)
            printf("%d\n", small.top());
        else if (sma_size > lar_size)
            printf("%d\n", small.top());
        else
            printf("%d\n", -large.top());
    }
    return 0;
}

사족

힙을 사용한다는 것을 알고 풀어서 풀 수 있었던 것 같아요