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

n은 최대 10만개까지 입력될 수 있기 때문에, 아래 코드처럼 정수가 하나씩 입력될 때마다 정렬하면 시간 초과가 발생한다.

#include <iostream>
#include <vector>
#include <algorithm> // std::sort
using namespace std;

int main() {
    cin.sync_with_stdio(0);
    cin.tie(0);

    // 정수를 하나씩 입력할 때마다 
    // 지금까지 입력된 수 중에서 '중간값'을 출력
    // 짝수개가 입력된 경우, 중간에 있는 두 수 중에서 작은 수 출력

    vector<int> v;

    int n;
    cin >> n;

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

        // 원소가 삽입될 때마다 오름차순 정렬
        sort(v.begin(), v.end());

        // 중간 인덱스에 해당하는 원소 출력
        int mid = v.size() / 2;
        if (v.size() % 2 == 0) { // 짝수개
            cout << min(v[mid], v[mid - 1]) << "\n";
        }
        else { // 홀수개 
            cout << v[mid] << "\n";
        }
    }

    return 0;
}

보다 효율적인 방법을 생각해야 하는데, 이 문제는 최대 힙과 최소 힙을 활용하여 중간값을 구할 수 있다.

알고리즘 이해

참고한 풀이: https://www.crocus.co.kr/625

항상 최대 힙의 루트 노드가 중간값이 될 수 있도록, 최대 힙과 최소 힙 사이의 균형을 맞추는 것이 핵심이다.

일단, 최대 힙의 루트 노드가 중간값이 되려면, 최대 힙의 노드 수는 최소 힙의 노드 수와 같거나 하나 더 많아야 한다. 짝수개의 정수가 입력된 경우, 중간에 있는 두 수 중에서 더 작은 값을 출력하므로, 최대 힙이 최소 힙보다 노드 수가 딱 하나만 더 많아야 한다. 최대 힙과 최소 힙의 노드 수가 2개 이상 차이가 나버리면, 최대 힙의 루트 노드는 중간값이 될 수 없다.

(최대 힙의 잎 노드 < 최대 힙의 루트 노드) < (최소 힙의 루트 노드 < 최소 힙의 잎 노드)

그리고 이 순서대로 값이 오름차순으로 커지려면, 최대 힙의 모든 노드는 최소 힙의 루트 노드보다 값이 작아야 한다. 만약 최대 힙의 루트 노드가 최소 힙의 루트 노드보다 값이 크면, 곧바로 swap해서 오름차순이 될 수 있게 한다.

최대 힙과 최소 힙으로 중간값을 구하는 알고리즘을 요약하면 다음과 같다.

1. 최대 힙의 노드 수는 최소 힙의 노드 수와 같거나 하나 더 많다.
2. 최대 힙의 모든 노드 <= 최소 힙의 루트 노드
   (if, 최대 힙의 루트 노드 > 최소 힙의 루트 노드 → 스왑)

위 규칙을 만족시키면서 노드를 하나씩 추가하면, 항상 최대 힙의 루트 노드가 전체 노드 중에서 중간값이 된다.

글만 읽었을 때는 이해하기가 어렵기 때문에, 최대 힙과 최소 힙을 직접 그려보면서 다시 이해해보자! 핵심은 최대 힙의 루트 노드가 중간값이 될 수 있도록 최대 힙과 최소 힙 사이의 균형을 맞춰준다는 것이다!

입력: 1 5 2 10 -99 7 5

최대 힙의 노드 수가 최소 힙의 노드 수와 같거나 하나 더 많아야 하므로, 최대 힙에 노드를 추가한다.

최대 힙과 최소 힙의 노드 수가 2개 이상 차이가 나면 안 되기 때문에 이번에는 최소 힙에 노드를 추가한다.

2를 최대 힙에 추가하면, 최대 힙 규칙에 따라 부모 노드에 가장 큰 값이 온다.

이미 최대 힙의 노드 수가 1개 더 많기 때문에 10은 최소 힙에 추가한다.

-99 추가

7 추가

5 추가

그래서 결과적으로, 1 1 2 2 2 2 5 가 중간값으로 출력된다.

코드로 구현

#include <iostream>
#include <queue>
using namespace std;

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

	priority_queue<int, vector<int>, less<int>> maxHeap;
	priority_queue<int, vector<int>, greater<int>> minHeap;

	int n;
	cin >> n;

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

        // 처음에 값이 없는 경우
        if (maxHeap.empty()) {
            maxHeap.push(x);
        }
        else {
            // 최대 힙의 크기가 더 크다면, 최소 힙에 값을 넣는다.
            if (maxHeap.size() > minHeap.size()) {
                minHeap.push(x);
            }
            else { // 크기가 같다면 최대 힙에 넣는다.
                maxHeap.push(x);
            }

            // 최대 힙의 루트 노드 > 최소 힙의 루트 노드 -> swap
            if (maxHeap.top() > minHeap.top()) {
                int maxTop = maxHeap.top();
                int minTop = minHeap.top();

                maxHeap.pop();
                minHeap.pop();

                maxHeap.push(minTop);
                minHeap.push(maxTop);
            }

        }

        // 최대 힙의 루트 노드가 중간값
        cout << maxHeap.top() << '\n';
	}

	return 0;
}