백준 2696번 - 중앙값 구하기

 

2696번: 중앙값 구하기

문제

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

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

입력

첫째 줄에 테스트 케이스의 개수 T(1 ≤ T ≤ 1,000)가 주어진다. 각 테스트 케이스의 첫째 줄에는 수열의 크기 M(1 ≤ M ≤ 9999, M은 홀수)이 주어지고, 그 다음 줄부터 이 수열의 원소가 차례대로 주어진다. 원소는 한 줄에 10개씩 나누어져있고, 32비트 부호있는 정수이다.

출력

각 테스트 케이스에 대해 첫째 줄에 출력하는 중앙값의 개수를 출력하고, 둘째 줄에는 홀수 번째 수를 읽을 때 마다 구한 중앙값을 차례대로 공백으로 구분하여 출력한다. 이때, 한 줄에 10개씩 출력해야 한다.

#include <iostream>
#include <queue>
#include <functional>

using namespace std;

int main() {
	int T;
	cin >> T;
	for (; T > 0; T--) {
		int N;
		cin >> N;
		cout << (N + 1) / 2 << endl;

		priority_queue<int, vector<int>, less<int>> small;
		priority_queue<int, vector<int>, greater<int>> big;

		int num;
		for (int i = 1; i <= N; i++) {
			cin >> num;
			if (small.empty() && big.empty()) {
				big.push(num);
			}
			else if (big.size() > small.size()) {
				if (num > big.top()) {
					small.push(big.top());
					big.pop();
					big.push(num);
				}
				else {
					small.push(num);
				}
			}
			else if (big.size() == small.size()) {
				if (num < small.top()) {
					big.push(small.top());
					small.pop();
					small.push(num);
				}
				else {
					big.push(num);
				}
			}

			if (i % 2 == 1) {
				cout << big.top() << ' ';
			}
		}
		cout << endl;
	}
	

}

우선순위 큐를 사용하는 문제. 우선순위 큐를 이 문제를 통해 처음알았다. 힙을 이용해 구현되어있어 FIFO가 아닌 우선순위에 따라 먼저 나갈 요소가 정해지는 형태의 큐이다.

 

1. 큰 수부터 먼저 나가는 우선순위 큐 small에는 중앙값보다 작은 수를, 작은 수부터 먼저 나가는 우선순위 큐 big에는 중앙값이상의 수를 넣어 big의 top이 중앙값이 될 수 있도록 만들어준다.

2. 이를 위해 big과 small의 크기를 맞춰줘야하므로 big이 small보다 큰데 입력된 요소가 big에 들어와야 한다면 big의 top을 small로 옮겨주어 개수를 맞춰주고 small에 들어가야 한다면 그냥 small에 넣어주면 big과 small의 개수가 같아진다.

big과 small의 크기가 같은 경우는 위와 반대로 해주면 big과 small의 크기는 같거나 big이 1 더 큰 경우 딱 두가지 밖에 나오지 않는다. 따라서 small이 big보다 큰 경우는 고려하지 않는다.

3. i가 홀수라면 중앙값인 big.top()을 출력하고 이를 테스트케이스만큼 반복한다.