가운데를 말해요 (백준1655)

문제

백준이는 동생에게 "가운데를 말해요" 게임을 가르쳐주고 있다. 백준이가 정수를 하나씩 외칠때마다 동생은 지금까지 백준이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 백준이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 백준이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 백준이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

첫째 줄에는 백준이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 백준이가 외치는 정수가 차례대로 주어진다. 정수는 -10,000보다 크거나 같고, 10,000보다 작거나 같다.

 

해결

코딩 테스트를 위한 자료 구조와 알고리즘 with C++ 이라는 책의 예제에서 만난 문제인데 어쩐지 코테 문제일 것 같더라니 역시나 백준에 비슷한 문제가 있어서 책 풀이를 정리해보았다.

단순히 정렬 후 가운데 값을 인덱스로 가져와도 동작은 하지만 시간 초과가 나버린다.

이 문제는 우선순위 큐 두 개를 사용하면 해결할 수 있다.

우선 maxHeap(root노드, 즉 [0]에 항상 최대값이 들어있는 힙)과 minHeap([0]에 항상 최소값이 들어있는 힙)을 생성한다.

maxHeap에는 중앙값보다 작은 수들이, minHeap에는 중앙값보다 큰 수들이 들어갈 예정으로, 중앙값은 maxHeap의 루트노드(중앙값보다 작은 값들 중에 가장 큰 값) 또는 minHeap의 루트 노트(중앙값보다 큰 값들 중에 가장 작은 값)가 될 것이다.

그리고 첫 번째 입력은 maxHeap에 넣는다.

현재 중앙값은 1이다.

두 번째 입력부터는 각 힙의 크기와 현재 중앙값을 고려하여 어느 힙에 넣을 것인지를 판단해야한다.

우선 minHeap이 maxHeap보다 작으므로 minHeap에 넣는 것을 먼저 고려해야 한다.

현재 중앙값은 1이다.(전체 수가 짝수개일 경우 둘 중 작은 것을 중앙값으로 하기 때문에)

이제 maxHeap과 minHeap의 크기가 동일하다.

따라서 다음 입력은 넣고자 하는 값과 현재 중앙값을 비교해서 입력해야한다.

넣고자 하는 값이 현재 중앙값보다 크면 minHeap에, 아니라면 maxHeap에 넣는다.

넣어진 값은 자동으로 힙 정렬이 된다.

현재 중앙값은 2이다.

이제 다시 maxHeap이 minHeap보다 사이즈가 작으므로 다음 입력을 maxHeap에 넣는 것을 고려해야 한다.

하지만 무조건적으로 수가 적은 Heap에 데이터를 넣으면 다음과 같은 문제가 생긴다.

maxHeap에 넣은 10이 힙 정렬에 의해 루트 노드가 되었는데, 이 루트 노드는 현재 중앙값인 minHeap의 루트 노드보다 큰 상태라서, maxHeap에는 중앙값보다 작은 수가 들어있어야 한다는 규칙을 위반한다.

maxHeap이 중앙값(2) 보다 큰 수 10을 들고 있다 -> 오류!

따라서 두 힙 간의 사이즈 차이가 있을 경우, 넣고자 하는 데이터와 현재 중앙값을 비교하여 데이터를 주고받는 과정이 필요하다.

예를 들어 maxHeap의 수가 적어 데이터를 넣고자 하는데 이 데이터가 중앙값(현재 minHeap의 루트 노드)보다 커서 규칙이 깨질 경우, 우선 minHeap의 루트 노드 데이터를 maxHeap에 넘긴다.

해당 값을 받은 maxHeap은 힙 정렬되어 해당 값이 루트 노드가 되도록 정렬할 것이다.

minHeap 또한 넘긴 데이터를 제외하고 다시 데이터를 정렬할 것이다.

이제 maxHeap보다 minHeap의 사이즈가 작으므로 minHeap에 넣고자 했던 데이터를 삽입한다.

minHeap이 힙 정렬을 하고 나면 다음과 같은 상태가 된다.

현재 중앙값은 2이다.

이제 다시 minHeap과 maxHeap의 사이즈가 같아졌다.

다음 입력은 아까와 마찬가지로 현재 중앙값보다 값이 크면 minHeap에, 값이 작으면 maxHeap에 넣는다.

현재 중앙값은 2이다.

minHeap이 maxHeap보다 사이즈가 작다.

이때 넣고자 하는 데이터 값 7이 현재 중앙값보다 크므로 minHeap에 넣어도 규칙이 깨지지 않는다.

(Heap은 인덱스 순서대로 정렬하지 않기 때문에 [0] 이후의 데이터들이 숫자 크기 순이 아닐 수 있지만, 루트 노드의 성질은 항상 성립한다)

현재 중앙값은 2이다.

이제 maxHeap과 minHeap의 사이즈가 같으므로 중앙값과 데이터를 비교하여 넣는다.

처음 코드가 실행될 때 maxHeap부터 입력했으므로 중앙값과 데이터가 같을 때는 maxHeap에 넣도록 한다.

최종 중앙값은 5이다.

이런 방식을 어떻게 생각해낸건지 처음 생각한 사람은 천재임이 분명하다.

 

코드 구현

#include <iostream>
#include <queue>
#include <vector>

struct median
{
	//최대 힙과 최소 힙을 생성한다
	std::priority_queue<int> maxHeap;
	std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap;

	void insert(int data)
	{
    	//처음 시작할 때는 maxHeap에 데이터를 넣는다
		if (maxHeap.size() == 0)
		{
			maxHeap.push(data);
			return;
		}
		
        //두 Heap의 사이즈가 같을 경우 입력하고자 하는 데이터와 중앙값을 비교하여
        //데이터가 중앙값보다 크다면 minHeap에, 아니라면 maxHeap에 넣는다
		if (maxHeap.size() == minHeap.size())
		{
			if (data <= get())
				maxHeap.push(data);
			else
				minHeap.push(data);

			return;
		}

		//두 힙 간의 사이즈 차이가 날 경우 데이터와 중앙값을 비교한다
        //maxHeap에 데이터를 넣으려는데 그 데이터가 중앙값(minHeap의 루트 노드)보다 클 경우
        //중앙값을 maxHeap에 옮겨준 후에 minHeap에 새 데이터를 넣는다.
		if (maxHeap.size() < minHeap.size())
		{
			if (data > get())
			{
				maxHeap.push(minHeap.top());
				minHeap.pop();
				minHeap.push(data);
			}
			else
				maxHeap.push(data);

			return;
		}

		//minHeap에 데이터를 넣으려는데 그 데이터가 중앙값(maxHeap의 루트 노드)보다 작을 경우
        //중앙값을 minHeap에 옮겨준 후에 maxHeap에 새 데이터를 넣는다.
		if (data < get())
		{
			minHeap.push(maxHeap.top());
			maxHeap.pop();
			maxHeap.push(data);
		}
		else
			minHeap.push(data);
	}

	//중앙값을 구한다
    //만약 중앙값이 두 개라면 둘 중 작은 값을 중앙값으로 한다
	int get()
	{
		if (maxHeap.size() == minHeap.size())
			return maxHeap.top() > minHeap.top() ? minHeap.top() : maxHeap.top();

		if (maxHeap.size() < minHeap.size())
			return minHeap.top();

		return maxHeap.top();
	}
};

int main()
{
   std::cin.tie(NULL);
    std::ios::sync_with_stdio(false);
    
    int nNum;
    std::cin >>nNum;
    
	median med;
    
    for(int i = 0; i < nNum ; i++)
    {
        int num;
        std::cin >> num;
        med.insert(num);
        std::cout << med.get() << '\n';
    }
    
	return 0;
}

 

테스트

테스트 결과는 백준 채점 결과로 대신한다.

푼 사람들 중에서도 상당히 빠른 결과인 것 같다.