파이썬 라이브러리 heapq

kkado·2022년 6월 9일

파이썬

목록 보기

2/4

원소를 정렬된 채로 저장할 수 있는 파이썬 내장 라이브러리 heapq에 대해 다룬다.

heap

먼저 heap 자료구조에 대해 간략히 알 필요가 있다.

min heap이란 트리 구조 안에서 부모 노드가 항상 자식 노드보다 작다 라는 특성이 있다.
반대로 max heap은 트리 구조 안에서 부모 노드가 항상 자식 노드보다 크다 라는 특성이 있다.

즉, min heap의 root 노드에는 그 트리의 최솟값, max heap의 root 노드에는 그 트리의 최댓값이 들어있음은 자명하다.

나중에 이진 트리와 힙을 다룰 때 더 알아보도록 하고, 본 내용은 라이브러리 소개하는 내용이니까 여기까지 ...

import 및 사용

import heapq

heapq 라이브러리 하나만 import 해주면 된다.

heapq는 별개의 자료구조가 아니라서 그냥 리스트를 사용하고 heapq 함수에 인자로 넘겨주는 방식을 사용한다.

즉 그냥 heap = []로 초기 생성 해주면 된다.

힙에 원소 추가

heapq.heappush() 함수를 사용한다. 넣을 리스트, 넣을 값을 차례로 인자로 넣어주면 된다.

heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 4)
heapq.heappush(heap, 2)

print(heap)

[1, 2, 4, 3]

순서를 아무렇게나 넣어도 자동으로 min heap이 되게끔 들어가는 것을 볼 수 있다.
그리고 두 번째 인자로 튜플을 넘겨줄 경우 그 튜플의 가장 첫번째 값을 기준으로 min heap이 생성된다.

내부적으로 heappush 함수는 O(logn)의 시간 복잡도를 가진다.

힙에 원소 삭제

heapq.heappop() 함수를 사용한다. 어떤 리스트에서 pop할지만 전달해 주면 된다. 그리고 삭제된 값, 즉 최솟값을 리턴한다.

print(heapq.heappop(heap))
print(heap)

1
[2, 3, 4]

이진 트리에서 원소를 삭제하는 heappop 함수도 O(logn)의 시간 복잡도를 가진다.

리스트를 힙으로 만들기

이미 원소가 차 있는 리스트도 힙으로 만들 수 있다. heapify 함수를 사용하면 된다.

heap = [1, 5, 2, 7, 4, 3, 9]

heapq.heapify(heap)
print(heap)

[1, 4, 2, 7, 5, 3, 9]

min heap의 조건을 만족하게끔 리스트가 재배열되었다.

최대 힙 (max heap 만들기)

heapq 함수들은 min heap을 기반으로 동작하기 때문에 약간의 꼼수를 써야한다.

heapq.heappush 함수의 두 번째 인자에 튜플을 넘겨주면 그 튜플의 첫번째 값을 기준으로 힙을 만든다는 점을 이용해서, 값에 음수를 취한 값을 넘겨주면 가장 큰 값이 가장 작은 값이 될 것이다.

즉 (-num, num) 의 튜플을 넘겨주면 되고, 나중에 값을 취할 때는 인덱스 1인 부분을 보면 되겠다.

heap = []

heapq.heappush(heap, (-5, 5))
heapq.heappush(heap, (-2, 2))
heapq.heappush(heap, (-3, 3))
heapq.heappush(heap, (-1, 1))
heapq.heappush(heap, (-7, 7))
heapq.heappush(heap, (-6, 6))
heapq.heappush(heap, (-4, 4))
heapq.heappush(heap, (-8, 8))

for i in heap :
    print(i[1])