문제 풀이

 

힙 2개를 이용하면 쉽게 해결할 수 있다.

각 각의 최대 힙과 최소 힙을 준비하자. 그리고 규칙에 따라 들어오는 input을 처리하면 된다.

 

1. max-heap 크기 > min-heap 크기

만족하면 min-heap에 넣는다.

아니라면 max-heap에 넣는다.

 

2. max-heap의 루트 노드 > min-heap의 루트 노드

만족하면 서로 swap 한다.

 

3. 홀수 번째 수를 읽을 때마다 중앙값을 처리해야 하므로,

홀수 번째이면 그때 max-heap의 루트 노드를 반환하면 된다.

 

문제의 예시를 보면 다음과 같이 흘러간다.

 

1이 들어오면

max-heap 1                
min-heap                  

 

5가 들어오면

max-heap 1                
min-heap 5                

 

4가 들어오면

max-heap 4 1              
min-heap 5                

 

3이 들어오면

max-heap 4 1              
min-heap 3 5              

규칙 2에 위반하므로, 서로 Swap 해준다.

max-heap 3 1              
min-heap 4 5              

2가 들어오면

max-heap 3 2 1            
min-heap 4 5              

빨간색으로 표시한 부분이 홀수 번째를 읽었을 때 중앙값이다.

 

소스 코드
#include <iostream>
#include <queue>
#include <vector>
using namespace std;

int main()
{
	ios_base::sync_with_stdio(0); cin.tie(0);
	int tc;
	cin >> tc;
	while (tc--)
	{
		int m;
		cin >> m;

		priority_queue<int> pq_max;
		priority_queue<int, vector<int>, greater<int>> pq_min;
		vector<int> ret;
		for (int i = 1; i <= m; i++)
		{
			int x;
			cin >> x;

			if (pq_max.size() > pq_min.size())
				pq_min.push(x);
			else
				pq_max.push(x);

			if (!pq_max.empty() && !pq_min.empty())
			{
				if (pq_max.top() > pq_min.top())
				{
					int a = pq_max.top(); pq_max.pop();
					int b = pq_min.top(); pq_min.pop();
					pq_max.push(b);
					pq_min.push(a);
				}
			}

			if (i % 2 == 1)
			{
				ret.push_back(pq_max.top());
			}
		}

		cout << ret.size() << '\n';
		for (int i = 0; i < ret.size(); i++)
		{
			if ((i + 1) % 10 == 0)
				cout << ret[i] << '\n';
			else
				cout << ret[i] << ' ';
		}
		cout << '\n';
	}
	return 0;
}

+ Recent posts