[Python] heapq 모듈

heapq 모듈

힙 자료구조

heapq모듈은 이진 트리(binary tree)기반의 최소 힙(min heap)자료구조 제공.

min heap을 사용하면 원소들이 항상 정렬된 상태로 삽입, 삭제되며 min heap에서 가장 작은 값은 언제나 인덱스 0, 즉 이진트리의 루트에 위치.
heap 자료구조를 이용해 데이터를 정렬하려면 heap[0]를 루프를 돌아 heappop() 해주면 된다.

 

모듈 임포트

import heapq

 

최소 힙 (min heap) 생성

heap = []

 

별개의 자료구조가 아닌 리스트를 힙처럼 다룰 수 있도록 하는 것

 

힙에 원소 추가 - heappush()

heap모듈의 heappush()함수를 이용하여 원소를 추가 할 수 있다.
첫번째 인자는 원소를 추가할 대상 리스트이며 두번빼 인자는 추가할 원소를 넘긴다.

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

 

힙에서 원소 삭제 - heappop()

heap모듈의 heappop()함수를 이용하여 힙에서 원소를 삭제 할 수 있다.
원소를 삭제할 대상을 인자로 넘기면, 인덱스 0에 있는 원소를 삭제 후에 그 값을 리턴한다.

heapq.heappop(heap)
print(heap) # [2, 4, 3, 5]

 

기존 리스트를 힙으로 변환

heap모듈의 heapify()를 사용하면 기존의 리스트를 리스트 힙으로 만들 수 있다.

heap = [5, 3, 4, 2, 1]
heapq.heapify(heap)
print(heap) # [1, 2, 4, 5, 3]

 

[응용] 최대 힙, 우선순위 큐

heapq모듈은 기본이 min heap인데, max heap을 구하려면 튜플(tuple)을 이용하면 된다.
이렇게 우선순위를 정하여 구할 수 있다.

num = [5, 3, 4, 2, 1]
heap = []

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

# 우선순위와 함께 출력
while heap:
    print(heapq.heappop(heap))

"""
(-5, 5)
(-4, 4)
(-3, 3)
(-2, 2)
(-1, 1)
"""

# 우선순위에 따라 정렬된 값만 출력
while heap:
    print(heapq.heappop(heap)[1])

"""
5
4
3
2
1
"""

 

시간 복잡도

메소드	시간복잡도
heappush()	O(logN)
heapop()	O(logN)
heapify()	O(N)

 

Reference

[파이썬] heapq 모듈 사용법 | Engineering Blog by Dale Seo