[python] heapq 모듈 사용해보기

오늘은 프로그래머스 문제를 풀면서 heap 자료구조를 사용해야하는 문제를 마주하면서 처음으로 heapq 모듈을 사용해보았다.

앞으로 코딩테스트를 준비하는데 있어 heapq모듈은 자주 쓰일 가능성이 있어 정리해보려고 한다.

 

 

 

heapq 모듈은 파이썬 표준 라이브러리에서 제공해주는 것으로 우선순위 큐 알고리즘이라고도 하는 힙(heap) 자료구조 알고리즘의 구현을 제공한다.

 

간단하게 힙(heap)이 무엇인지 알아보고 넘어가자.

힙(heap) 자료구조란 완전 이진 트리의 일종으로 우선순위 큐를 위하여 만들어진 자료구조이다.

자료구조 특성 상 여러 개의 값들 중에서 최댓값이나 최솟값을 빠르게 찾아낼 수 있다.
힙의 형태는 부모 노드의 키 값이 자식 노드의 키 값보다 항상 큰 혹은 작은 이진 트리로 힙에서는 중복된 값을 허용한다.

힙은 최대 힙(max heap)과 최소 힙(min heap)으로 나눌 수 있다.
최대 힙은 부모 노드의 키 값이 자식 노드의 키 값보다 크거나 같은 완전 이진 트리를 말하고
최소 힙은 부모 노드의 키 값이 자식 노드의 키 값보다 작거나 같은 완전 이진 트리를 말한다.

최대 힙과 최소 힙의 트리 구조를 그림으로 보면 아래와 같다.

출처: https://www.geeksforgeeks.org/heap-data-structure/minheapandmaxheap/

 

heapq 모듈에서는 우리가 직접 힙 알고리즘을 구현할 필요없이 여러 함수를 이용하여 힙을 쓸 수 있다.

그러면 heapq 모듈에 있는 함수는 무엇이 있는지 알아보자.

1) heapq.heappush(heap, item)

힙 불변성을 유지하면서, item 값을 heap으로 푸시한다.

import heapq

heap = []
heapq.heappush(heap, 4)
heapq.heappush(heap, 3)
heapq.heappush(heap, 2)
heapq.heappush(heap, 5)
heapq.heappush(heap, 1)

print(heap)
# 결과: [1, 2, 3, 5, 4]

결과를 보면 알겠지만 기본적으로 heapq 모듈에서 제공하는 heap은 최소 힙구조이다.

최대 힙을 구현하기 위한 방법은 아래에서 알아보자.

 

 

2) heapq.heappop(heap)

힙 불변성을 유지하면서, heap에서 가장 작은 값을 팝하고 반환한다. 만약에 힙이 비어 있으면 IndexError가 발생한다. 팝이 아닌 단순히 heap에서 가장 작은 값을 원한다면 heap[0]을 사용하면된다.

기본적으로 heapq 모듈에서 제공하는 heap은 최소 힙구조이기때문에 설명에서 최소라고 한 것이다.

만약, 최대 힙으로 구현했다면 최대 값이 나올 것이다.

import heapq

heap = []
heapq.heappush(heap, 4)
heapq.heappush(heap, 3)
heapq.heappush(heap, 2)
heapq.heappush(heap, 5)
heapq.heappush(heap, 1)

min_value = heapq.heappop(heap)
print(min_value)
# 결과: 1

min_value = heapq.heappop(heap)
print(min_value)
# 결과: 2

 

 

3) heapq.heapify(x)

heapify는 기존에 존재하는 리스트를 최소 힙구조로 만들어주는 기능을 한다.

주의해야할 점은 return이 없다는 것이다. 원래 존재하던 리스트가 최소 힙구조로 변하는 것이다.

import heapq

arr = [1, 2, 6, 5, 4, 3]
heapq.heapify(arr)

print(arr)
#결과: [1, 2, 3, 5, 4, 6]


# heapq.heapify()이 반환값이 없다는 것을 알아보기 위한 실습
arr2 = [1, 5, 3, 4, 9, 7]
heap = heapq.heapify(arr2)

print(arr)
#결과: None

 

 

4) heapq.heappushpop(heap, item)

힙에 키 값을 푸시하고 힙에 있는 값중 가장 작은 값을 반환해주는 함수이다.

만약에 heappush()한 다음 바로 heappop()을 한다면 이 함수를 사용하는 것이 효율적이다.

import heapq

heap = []
heapq.heappush(heap, 4)
heapq.heappush(heap, 3)
heapq.heappush(heap, 2)
heapq.heappush(heap, 5)
heapq.heappush(heap, 1)

min_value = heapq.heappushpop(heap, 7)
print(heap, min_value)

# 결과: [2, 4, 3, 5, 7], 1

 

 

5) heapq.heapreplace(heap, item)

heapq.heappushpop()함수와 유사하지만 실행되는 순서가 다른 함수이다.

heap에서 가장 작은 값을 팝하여 반환한 후에 새로운 값을 푸시한다. 만약에 힙이 비어 있으면 IndexError가 발생한다.

이 함수는  heappop()한 다음 heappush()를 쓰는 것보다 더 효율적이며 고정 크기 힙을 사용할 때 더 적합 할 수 있다. 

여기서 주의해야할 점은 반환된 값(팝)이 이후에 푸시된 item보다 클 수 있다는 것이다.

import heapq

heap = []
heapq.heappush(heap, 4)
heapq.heappush(heap, 3)
heapq.heappush(heap, 2)
heapq.heappush(heap, 5)
heapq.heappush(heap, 7)

min_value = heapq.heapreplace(heap, 1)
print(heap, min_value)

# 결과: [1, 4, 3, 5, 7], 2

 

 

 

지금까지 heapq 모듈에서 제공하는 함수를 알아보았다.

(위에서 언급한 함수가 전부가 아니므로 더 알고 싶다면 아래 참고 자료 링크 두 번째를 참고하면 좋을 것같다.)

그러면 이제 heapq 모듈을 이용해서 최대 힙을 구현해보는 것을 알아보자.

최대 힙(max heap)을 구현하기 위해서는 힙에 튜플(tuple)을 값으로 푸시한거나 팝하면 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 것을 이용하면 된다.

우리가 push하고자하는 각 값에 대한 우선 순위를 구한 후 (우선 순위, 값) 구조의 튜플(tuple)을 힙에 푸시하거나 팝하면 된다.

만약 최대값을 팝하고 싶다면 heapq.heappop(heap)을 하면 튜플이 나오므로 튜플에서 인덱스 1에 있는 값을 얻으면 된다.

 

글로만 보면 이해가 잘 안되므로 코드를 보면서 알아보자.

import heapq

arr = [4, 1, 7, 3, 8, 5]
min_heap = []
max_heap = []

# 최소힙 만들기
heapq.heapify(arr)
min_heap = arr


# 최대힙 만들기
for key in arr:
  heapq.heappush(max_heap, (-key, key))

print(min_heap)
#결과 : [1, 3, 5, 4, 8, 7]
print(max_heap)
#결과 : [(-8, 8), (-5, 5), (-7, 7), (-1, 1), (-4, 4), (-3, 3)]

기본적으로 heapq.heapify()를 이용하면 최소 힙이 된다는 것은 알고 있다.

최대 힙을 만들기 위해 각각의 값의 음수 형태를 우선순위로 정하여 (우선순위, 값) 튜플을 푸시한다면 우선순위 기준으로 최소 힙이 만들어지기 때문에 결과적으로는 우리가 원하는 값에 대해서는 최대 힙이 만들어지는 것이다.

 

 

---

<참고 자료>

https://gmlwjd9405.github.io/2018/05/10/data-structure-heap.html 

https://docs.python.org/ko/3/library/heapq.html

https://www.daleseo.com/python-heapq/