백준 2696: 중앙값 구하기

문제 요약

어떤 수열을 읽고, 홀수번째 수를 읽을 때 마다, 지금까지 입력받은 값의 중앙값을 출력하는 프로그램을 작성하시오.

예를 들어, 수열이 1, 5, 4, 3, 2 이면, 홀수번째 수는 1번째 수, 3번째 수, 5번째 수이고, 1번째 수를 읽었을 때 중앙값은 1, 3번째 수를 읽었을 때는 4, 5번째 수를 읽었을 때는 3이다.

문제 분류

• 자료 구조
• 우선순위 큐

문제 풀이

중앙값을 root로 잡고, 양측으로 두 개의 우선순위 큐를 가지면 된다. root보다 작은 값만을 원소로 가지는 큐(q1)와 root보다 큰 값만을 원소로 가지는 큐(q2)이다. 이때 q1은 최대 힙으로, q2는 최소 힙으로 만든다. 이렇게 골격을 짰다면, 나머지는 양 측의 큐의 크기를 동일하도록 보장하면서 입력값에 대해 처리해주면 된다. 또한, 각 큐의 top은 중앙값의 후보가 된다.

먼저, 첫번째 원소가 입력 되었을 경우에는, 무조건 자기 자신이 root가 된다. 그 이후 입력된 원소마다 다음의 연산을 하면 된다.

입력된 원소를 in이라고 하면, 우선 q1과 q2의 크기 및 모든 경우를 비교하면서 해야할 연산을 결정한다. 그 경우는 각각 다음과 같다.

1. q1의 크기가 q2의 크기보다 작다면~ (q1에 원소를 삽입해야 함을 의미)

   1. in이 q2의 top보다 크다면
      • 다음의 경우는 in이 당연히 root보다도 크며, 후보가 될 수도 없다. root를 q1에 삽입하고, q2의 top을 root로 사용하며, q2에 in을 삽입한다.
   2. 그렇지 않고 단순히 in이 root보다 크다면
      • in이 q2의 top보다는 작으므로 중앙값으로 판단된다. root를 q1에 삽입하고, in을 root로 사용한다.
   3. in이 root보다 작을 경우
      • in이 root보다 작은 q1에 속하게 되므로, q1에 in을 삽입한다.

2. 그렇지 않다면~(q2에 원소를 삽입해야 함을 의미)

   1. n이 q1의 top보다 작은 경우
      • in은 q1에 속하게 되므로, q2에 root를 삽입하고, q1의 top을 root로 사용하며, q1에 in을 삽입한다.
   2. in이 root보다 작은 경우
      • in이 중앙값이 판단되므로 root를 q2에 삽입하고, in을 root로 사용한다.
   3. in이 root보다 큰 경우
      • in이 root보다 큰 q2에 속하게 되므로, q2에 in을 삽입한다.

다음의 연산이 가능한 것은 q1은 최대 힙, q2를 최소 힙으로 구성하는 것에 있다. q1의 값들은 모두 q1의 top보다는 작을 것이고, q2의 값들 역시 q2의 top보다는 클 것이다. 중앙값보다 작은 값들 중 가장 큰 값, 중앙값보다 큰 값들 중 가장 작은 값만 알면 된다. 각각의 top이 곧 다음 연산 시 root후보가 되는 것이다.

풀이 코드

#include <stdio.h>
#include <iostream>
#include <algorithm>
#include <queue>

using namespace std;

int main()
{
	int t, n, in, root;
	priority_queue<int> q1;
	priority_queue<int, vector<int>, greater<>> q2;
	vector<int> res;

	cin >> t;
	while (t--) {
		scanf("%d", &n);
		for (int i = 1; i <= n; i++) {
			scanf("%d", &in);
			if (i == 1)
				root = in;
			else {
				if (q1.size() < q2.size()) {
					if (q2.top() < in) {
						q1.push(root);
						root = q2.top();
						q2.pop();
						q2.push(in);
					}
					else if (root < in) {
						q1.push(root);
						root = in;
					}
					else q1.push(in);
				}
				else {
					if (!q1.empty()) {
						if (q1.top() > in) {
							q2.push(root);
							root = q1.top();
							q1.pop();
							q1.push(in);
						}
						else if (root > in) {
							q2.push(root);
							root = in;
						}
						else q2.push(in);
					}
					else {
						if (root < in) q2.push(in);
						else q1.push(in);
					}
				}
			}
			if (i % 2)
				res.push_back(root);
		}
		printf("%d\n", res.size());
		for (int i = 0; i < res.size(); i++) {
			printf("%d ", res[i]);
			if (i % 10 == 9) printf("\n");
		}
		printf("\n");
		res.clear();
		q1 = priority_queue<int>();
		q2 = priority_queue<int, vector<int>, greater<>>();
	}
	return 0;
}