파이썬의 heapq 모듈로 힙 자료구조 사용하기

데이터를 정렬된 상태로 저장하기 위해서 사용하는 파이썬의 heapq(힙큐) 내장 모듈에 대해서 알아보겠습니다.

힙 자료구조

heapq 모듈은 이진 트리(binary tree) 기반의 최소 힙(min heap) 자료구조를 제공합니다. 자바에 익숙하신 분이라면 PriorityQueue 클래스를 생각하시면 이해가 쉬우실 것 같습니다.

min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, min heap에서 가장 작은값은 언제나 인덱스 0, 즉, 이진 트리의 루트에 위치합니다. 내부적으로 min heap 내의 모든 원소(k)는 항상 자식 원소들(2k+1, 2k+2) 보다 크기가 작거나 같도록 원소가 추가되고 삭제됩니다.

heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2]

예를 들어, 아래 그림은 위 공식을 만족시키는 간단한 min heap의 구조를 보여주고 있습니다.

     1  ---> root
   /   \
  3     5
 / \   /
4   8 7

모듈 임포트

우선 heapq 모듈은 내장 모듈이기 때문에 파이썬만 설치되어 있으면 다음과 같이 간단하게 임포트 후에 힙 관련 함수를 사용할 수 있습니다.

from heapq import heappush, heappop //, ... 다른 함수들

최소 힙 생성

heapq 모듈에은 파이썬의 보통 리스트를 마치 최소 힙처럼 다룰 수 있도록 도와줍니다. 자바의 PriorityQueue 클래스처럼 리스트와 별개의 자료구조가 아닌 점에 유의해야 합니다.

그렇게 때문에, 그냥 빈 리스트를 생성해놓은 다음 heapq 모듈의 함수를 호출할 때 마다 이 리스트를 인자로 넘겨야 합니다. 다시말해, 파이썬에서는 heapq 모듈을 통해서 원소를 추가하거나 삭제한 리스트가 그냥 최소 힙입니다.

heap = []

힙에 원소 추가

heapq 모듈의 heappush() 함수를 이용하여 힙에 원소를 추가할 수 있습니다. 첫번째 인자는 원소를 추가할 대상 리스트이며 두번째 인자는 추가할 원소를 넘깁니다.

from heapq import heappush

heappush(heap, 4)
heappush(heap, 1)
heappush(heap, 7)
heappush(heap, 3)
print(heap)

[1, 3, 7, 4]

가장 작은 1이 인덱스 0에 위치하며, 인덱스 1(= k)에 위치한 3은 인덱스 3(= 2k + 1)에 위치한 4보다 크므로 힙의 공식을 만족합니다. 내부적으로 이진 트리에 원소를 추가하는 heappush() 함수는 O(log(n))의 시간 복잡도를 가집니다.

힙에서 원소 삭제

heapq 모듈의 heappop() 함수를 이용하여 힙에서 원소를 삭제할 수 있습니다. 원소를 삭제할 대상 리스트를 인자로 넘기면, 가장 작은 원소를 삭제 후에 그 값을 리턴합니다.

from heapq import heappop

print(heappop(heap))
print(heap)

1
[3, 4, 7]

가장 작았던 1이 삭제되어 리턴되었고, 그 다음으로 작었던 3이 인덱스 0으로 올라왔습니다.

print(heappop(heap))
print(heap)

3
[4, 7]

가장 작었던 3이 삭제되어 리턴되었고, 그 다음으로 작았던 4가 인덱스 0으로 올라왔습니다. 내부적으로 이진 트리로 부터 원소를 삭제하는 heappop() 함수도 역시 O(log(n))의 시간 복잡도를 가집니다.

최소값 삭제하지 않고 얻기

힙에서 최소값을 삭제하지 않고 단순히 읽기만 하려면 일반적으로 리스트의 첫번째 원소에 접근하듯이 인덱스를 통해 접근하면 됩니다.

print(heap[0])

4

여기서 주의 사항은 인덱스 0에 가장 작은 원소가 있다고 해서, 인덱스 1에 두번째 작은 원소, 인덱스 2에 세번째 작은 원소가 있다는 보장은 없다는 것입니다. 왜냐하면 힙은 heappop() 함수를 호출하여 원소를 삭제할 때마다 이진 트리의 재배치를 통해 매번 새로운 최소값을 인덱스 0에 위치시키기 때문입니다.

따라서 두번째로 작은 원소를 얻으려면 바로 heap[1]으로 접근하면 안 되고, 반드시 heappop()을 통해 가장 작은 원소를 삭제 후에 heap[0]를 통해 새로운 최소값에 접근해야 합니다.

기존 리스트를 힙으로 변환

이미 원소가 들어있는 리스트 힙으로 만들려면 heapq 모듈의 heapify()라는 함수에 사용하면 됩니다.

from heapq import heapify

heap = [4, 1, 7, 3, 8, 5]
heapify(heap)
print(heap)

[1, 3, 5, 4, 8, 7]

heapify() 함수에 리스트를 인자로 넘기면 리스트 내부의 원소들의 위에서 다룬 힙 구조에 맞게 재배치되며 최소값이 0번째 인덱스에 위치됩니다. 즉, 비어있는 리스트를 생성한 후 heappush() 함수로 원소를 하나씩 추가한 효과가 납니다. heapify() 함수의 성능은 인자로 넘기는 리스트의 원소수에 비례합니다. 즉 O(n)의 시간 복잡도를 가집니다.

heapify() 함수를 사용할 때 주의할 점은 새로운 리스트를 반환하는 것이 아니라 인자로 넘긴 리스트에 직접 변경한다는 것입니다. 따라서 원본 리스트의 형태를 보존해야되는 경우에는 반드시 해당 리스트를 복제한 후에 인자로 넘겨야 하겠습니다.

nums = [4, 1, 7, 3, 8, 5]

heap = nums[:]
heapify(heap)

print(nums)
print(heap)

[4, 1, 7, 3, 8, 5]
[1, 3, 5, 4, 8, 7]

[응용] 최대 힙

heapq 모듈은 최소 힙(min heap)을 기능만을 동작하기 때문에 최대 힙(max heap)으로 활용하려면 약간의 요령이 필요합니다. 바로 힙에 튜플(tuple)를 원소로 추가하거나 삭제하면, 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 원리를 이용하는 것입니다.

따라서, 최대 힙을 만들려면 각 값에 대한 우선 순위를 구한 후, (우선 순위, 값) 구조의 튜플(tuple)을 힙에 추가하거나 삭제하면 됩니다. 그리고 힙에서 값을 읽어올 때는 각 튜플에서 인덱스 1에 있는 값을 취하면 됩니다. (우선 순위에는 관심이 없으므로 )

from heapq import heappush, heappop

nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
  heappush(heap, (-num, num))  # (우선 순위, 값)

while heap:
  print(heappop(heap)[1])  # index 1

8
7
5
4
3
1

[응용] n번째 최소값/최대값

최소 힙이나 최대 힙을 사용하면 n 번째로 작은 값이나 n 번째로 큰 값을 효과적으로 구할 수 있습니다.

n 번째 최소값을 구하기 위해서는 주어진 배열로 힙을 만든 후, heappop() 함수를 n 번 호출하면 됩니다.

from heapq import heappush, heappop

def nth_smallest(nums, n):
    heap = []
    for num in nums:
        heappush(heap, num)

    nth_min = None
    for _ in range(n):
        nth_min = heappop(heap)

    return nth_min

print(nth_smallest([4, 1, 7, 3, 8, 5], 3))

4

heapify() 함수를 활용하면 힙을 만들 때 굳이 루프를 돌면서 숫자를 매 번 하나씩 추가해줄 필요는 없겠죠?

from heapq import heapify, heappop

def nth_smallest(nums, n):
    heapify(nums)

    nth_min = None
    for _ in range(n):
        nth_min = heappop(nums)

    return nth_min

사실 heapq 모듈에 이미 이러한 용도로 사용할 수 있는 nsmallest()라는 함수가 존재합니다. nsmallest() 함수는 주어진 리스트에서 가장 작은 n개의 값을 담은 리스트를 반환하는데요. 그 결과 리스트의 마지막 값이 n 번째 작은 값이 되겠습니다.

from heapq import nsmallest

nsmallest(3, [4, 1, 7, 3, 8, 5])[-1]

반대로 n 번째 최대값을 구할 때는 nlargest() 함수를 사용하면 되겠습니다.

from heapq import nlargest

nlargest(3, [4, 1, 7, 3, 8, 5])[-1]

[응용] 힙 정렬

힙 정렬(heap sort)은 위에서 설명드린 힙 자료구조의 성질을 이용한 대표적인 정렬 알고리즘입니다.

from heapq import heappush, heappop

def heap_sort(nums):
  heap = []
  for num in nums:
    heappush(heap, num)

  sorted_nums = []
  while heap:
    sorted_nums.append(heappop(heap))
  return sorted_nums

print(heap_sort([4, 1, 7, 3, 8, 5]))

[1, 3, 4, 5, 7, 8]

전체 코드

본 포스팅에서 제가 작성한 전체 코드는 아래에서 직접 확인하고 실행해보실 수 있습니다.

https://dales.link/6aq

마치면서

지금까지 파이썬의 heapq 내장 모듈을 다양한 방법으로 사용해보았는데요. 사실 힙이 파이썬에서 리스트나 사전처럼 매일 사용하는 자료구조는 아니지만 그래도 잘 공부해두시면 나중에 분명히 유용하게 쓰실 때가 있으실 거라 생각합니다.