[백준 알고리즘] 1655번: 가운데를 말해요

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

문제

문제
수빈이는 동생에게 “가운데를 말해요” 게임을 가르쳐주고 있다. 수빈이가 정수를 하나씩 외칠때마다 동생은 지금까지 수빈이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 수빈이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 수빈이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 수빈이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

입력
첫째 줄에는 수빈이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 수빈이가 외치는 정수가 차례대로 주어진다. 정수는 -10,000보다 크거나 같고, 10,000보다 작거나 같다.

출력
한 줄에 하나씩 N줄에 걸쳐 수빈이의 동생이 말해야하는 수를 순서대로 출력한다.

풀이

힙소트를 계속 해주면서 중간 인덱스를 참조하는 문제인 줄 알았는데 완전히 낭패였다. 시간초과가 계속 나서 어떻게 해결해야할지 고민하다 답이 안나오는 것 같아 구글링으로 도움을 받았다. 이 문제는 우선순위 큐 두개를 사용해서 push 와 pop만을 사용해서 해결해야 했다. Heapsort는 Օ(nlogn) 으로 절대로 느리지는 않지만, 이 문제에는 적합하지 않다. push 와 pop 은 Օ(lg n)의 시간복잡도를 가지는데, 트리의 높이만큼 부모노드와 자식노드의 비교 및 교환을 수행하기 때문이다.

두 개의 우선순위 큐를 사용해서 문제를 해결하는 방법은 다음과 같다.

1. 최대 우선순위 큐, 최소 우선순위 큐 두 개를 준비한다.
2. 최소 우선순위 큐는 항상 최대 우선순위 큐 보다 큰 값을 가진다.
3. 최대 우선순위 큐는 항상 최소 우선순위 큐 보다 같거나 많은 노드를 가진다.

이 구조를 유지하게되면, 어떤 값이 들어왔을 때, 총 노드의 개수가 짝수일 때는 자연스럽게 최대 우선순위 큐에 값을 넣게 된다. 이 때 최댓값이 큐의 top의 위치하게 되기 때문에, 최소 우선순위 큐의 top에 있는 노드의 값과 비교해주어야할 필요가 생긴다. 왜냐하면 최대 우선순위 큐가 최소 우선순위 큐보다 하나 더 많은 노드를 갖게하는 것이 중간 값을 최대 큐에 저장하기 위해서인데, 만약 새로 들어온 값이 최소 큐의 top보다 크다면, 현재 최소 큐의 top에 있는 값이 중간값이 되기 때문이다. 이 구조를 유지하기 위해서 양쪽 노드를 조건에 따라 swap 하면서 최대 큐의 top을 계속 출력해주면 중간값을 지속적으로 얻을 수 있다.

코드

#include <cstdio>
#include <vector>
#include <queue>
#include <functional>

using namespace std;

int main (){
    int N;
    scanf("%d", &N);

    priority_queue<int, vector<int>, less<int>> pqMax;
    priority_queue<int, vector<int>, greater<int>> pqMin;

    while(N--){
        int n;
        scanf("%d", &n);
        if (pqMax.empty() || pqMax.size() == pqMin.size()){ // 무조건 최대 큐에 들어가야하는 경우
            pqMax.push(n);
        }
        else {
            pqMin.push(n);
        }

        if (!pqMax.empty() && !pqMin.empty()){
            if (pqMin.top() < pqMax.top()){ // 최대 큐는 항상 최소 작거나 같은 값을 가져야 한다.
                int tempMax = pqMax.top();
                int tempMin = pqMin.top();
                pqMax.pop();
                pqMin.pop();

                pqMax.push(tempMin);
                pqMin.push(tempMax);
            }
        }
        printf("%d\n", pqMax.top());
    }
}