BAEKJOON #1655) 가운데를 말해요(Python)

숫자를 입력받을 때마다 중앙값(숫자의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수)을 찾아서 출력하는 문제이다. 

처음에는 숫자를 입력받을 때마다 배열에 넣고 정렬하여 홀수면 n/2번째 인덱스를, 짝수면 (n-1)/2번째 인덱스를 출력하면 된다고 생각하여 정렬 알고리즘을 구현하여 제출했다.

뚜둥-! 퀵소트로 정렬하면 시간 초과가 나고, 병합정렬(merge sort)을 이용하면 메모리 초과가 난다😇

정렬 알고리즘을 이용하여 풀 수는 없다. 조금 찾아보니 힙을 사용하여 중앙값을 찾는 방법을 이용하면 된다고 한다.

Data Structure - Heap (힙) | Python 구현

힙은 최댓값과 최솟값을 찾을 수 있는 빠른 방법이므로 내가 찾고 싶은 중앙값을 힙으로 찾으려면 최대힙과 최소힙을 모두 만들어야 한다. 대충 이야기하면 작은 수들은 최대 힙에 넣고 큰 수들을 최소 힙에 넣으면,  최대힙에서의 최댓값 또는 최소힙에서의 최솟값이 중앙값이 될 것이다.

문제에서 숫자의 개수가 짝수일 때의 중앙값은 가운데 두 수 중 작은 수를 출력하라고 하였으므로, 우리는 두개로 나누어진 힙에서 항상 최대 힙에서의 최댓값을 출력하기로 한다.

 

1. 최대 힙에 넣을지, 최소 힙에 넣을지 결정

: 최대 힙의 루트 노드가 전체 숫자 중에서 중앙값이 되려면 항상 힙의 길이를 같거나 최대 힙의 길이를 하나만 더 크게 유지해야 한다.

• 숫자의 개수가 짝수일 때) 최대 힙의 노드 개수 = 최소 힙의 노드 개수
• 숫자의 개수가 홀수일 때) 최대 힙의 노드 개수 = 최소 힙의 노드 개수 + 1

2. 힙에 넣었을 때 최대 힙의 최상위 노드가 중앙값이 되도록 

: 최대 힙에 넣었든, 최소 힙에 넣었든 우리가 출력하기로한 중앙값의 위치의 값(즉, 최대힙의 루트노드의 값)이 최소 힙의 루트노드(최솟값)보다 크다면 값의 위치를 바꾸어주어야 한다. (최소 힙의 노드의 값들은 최대힙의 루트노드(중앙값)보다 항상 크거나 같아야 한다.)

from sys import stdin
import heapq


n = int(stdin.readline())
left_heap = list()  # max heap
right_heap = list()  # min heap

for i in range(n):
    val = int(stdin.readline())

    if len(left_heap) == len(right_heap):
        heapq.heappush(left_heap, [-val, val])
    else:
        heapq.heappush(right_heap, val)

    if right_heap and left_heap[0][1] > right_heap[0]:
        left_popped_val = heapq.heappop(left_heap)[1]
        right_popped_val = heapq.heappop(right_heap)
        heapq.heappush(left_heap, [-right_popped_val, right_popped_val])
        heapq.heappush(right_heap, left_popped_val)

    print(left_heap[0][1])

파이썬에서는 힙 자료구조를 제공해주는 heapq 라이브러리가 있기 때문에 따로 힙을 구현하지는 않았다. 계속 힙을 그리면서 알고리즘을 작성하다보니 left_heap, right_heap 이렇게 작성하였는데, left_heap이 작은 값들을 모아놓는 최대힙, right_heap이 큰 값들을 모아놓는 최소힙이다.

 

첫번째 if-else문이 들어온 숫자를 최대힙에 넣을지, 최소힙에 넣을지 결정하는 조건문이다. 즉, 각각의 힙의 개수를 맞춰주는 과정이다.

두번째 if문은 힙에 숫자를 추가하였을 때 현재의 최대 힙의 루트 노드가 중앙값이 아니라면 값의 위치를 바꿔주는 과정이다.