[백준] 1655번: 가운데를 말해요

📌 문제

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

📌 설명

수들을 순차적으로 받으면서 각 반복마다 수열의 중앙값을 출력하는 문제이다. 처음에 생각한 방법은 우선순위 큐나 set 자료구조를 통해 최대/최소값을 빠르게 찾은 다음, 선형적으로 탐색하여 중앙값을 찾는 방법이었다. 결과적으로 이 방법은 TLE이다. 현재 원소의 수를 $k$라고 하면, 원소를 삽입하는 데 $O(log(k))$, 중앙값을 찾는 데 $O(k)$의 시간복잡도를 가진다. 이를 $N$번 반복하므로 전체 시간복잡도는 $O(N^2)$이다.

이 문제를 해결하는 핵심 아이디어는 두 개의 우선순위 큐를 사용하는 것이다. 중앙값보다 작은 값들을 저장하는 내림차순 기반 우선순위 큐 left, 중앙값보다 큰 값들을 저장하는 오름차순 기반 우선순위 큐 right를 사용한다. left 또는 right의 크기가 지나지게 커지는 경우에 대한 처리 방법을 추가로 생각해야 한다. 결과적으로 이 방법은 $O(N log(N))$의 시간복잡도를 가져 제한 시간 내 문제를 해결할 수 있다.

📌 코드

#include <iostream>
#include <queue>

using namespace std;

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

	int n;
	cin >> n;

	priority_queue<int> left; // 중앙값을 기준으로 왼쪽에 위치한 수
	priority_queue<int, vector<int>, greater<int>> right; // 중앙값을 기준으로 오른쪽에 위치한 수
	for (int i = 0; i < n; i++) {
		int x;
		cin >> x;

		// right에서 중앙값에 가까운 수(right.top())이 x보다 작으면 단순히 right에 x 삽입
		if (!right.empty() && right.top() < x)
			right.push(x);
		else
			left.push(x);

		int mid = i / 2 + 1;

		// left에 수가 치우친 경우
		if (left.size() > mid) {
			right.push(left.top());
			left.pop();
		}
		else if (left.size() < mid) {
			left.push(right.top());
			right.pop();
		}

		cout << left.top() << "\n";
	}
}