[Python] heapq 모듈

힙 자료구조

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

heapq 모듈은 이진 트리 기반의 최소 힙(min heap) 자료구조를 제공합니다. 자바의 PriorityQueue 클래스와 비슷하다고 생각하시면 될 듯 합니다.

min heap에서 가장 작은 값은 언제나 0번 인덱스(이진 트리의 루트)에 위치합니다. min heap 내의 모든 원소(k)는 자식 원소(2k+1, 2k+2)보다 크기가 작거나 같도록 원소가 추가되고 삭제됩니다.

모듈 임포트

import heapq

내장 모듈이기 때문에 간단하게 임포트 할 수 있습니다.

힙 생성

heapq 모듈은 보통 리스트를 힙처럼 다룰 수 있도록 도와줍니다. 자바의 PriorityQueue 클래스 처럼 다른 자료구조가 아니라는 점에 유의해야 합니다.

1. 빈 리스트를 생성하여 힙처럼 다루는 방법

heap = []
heapq.heappush(heap, 4) # 힙에 원소 추가하는 함수
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3)
print(heap)

[1, 3, 7, 4]	# 실행 결과

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

2. 기존에 있던 리스트를 Heapify를 이용해 heap으로 만드는 방법

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

[1, 2, 7, 3, 4] # 실행 결과

heapify() 함수를 이용하면 리스트 내부의 원소들이 힙 구조에 맞게 재배치되며 최소값이 0번째 인덱스에 위치하게 됩니다. heapify() 함수의 성능은 인자로 넘기는 리스트의 원소수에 비례합니다. O(N) 의 시간 복잡도를 가집니다.

힙에서 원소 삭제

heappop()함수를 이용하여 힙에서 원소를 삭제할 수 있습니다. 원소를 삭제하고 싶은 리스트를 인자로 넘기면, 가장 적은 원소를 삭제한 후에 그 값을 리턴합니다.

heap = [1, 2, 7, 3, 4]
print(heapq.heappop(heap))
print(heap)

1
[2, 3, 7, 4]	# 실행 결과

가장 작았던 1 이 삭제되면서 리턴되었고, 그 다음으로 작은 수인 2 가 인덱스 0으로 왔습니다. 그리고 [2, 7, 3, 4] 리스트가 [2, 3, 7, 4] 리스트로 heap 구조로 정렬 되었습니다. heappop() 함수는 O(logN) 의 시간 복잡도를 가집니다.

출처 : heapq 모듈 사용법