최소 힙과 최대 힙 개념은 위 글에서 정리했다.
heapq
파이썬에서 제공하는 최소 힙 자료구조 모듈
자바의 PriorityQueue 클래스와 유사하다.
heapq를 사용하면 원소가 정렬된 상태로 추가되고 삭제된다.
모듈 임포트
내장 묘듈이기 때문에 별도 라이브러리 설치 없이 사용할 수 있다.
import heapq
최소 힙 생성
heapq는 파이썬의 일반 리스트를 최소 힙처럼 다룰 수 있도록 한다.
빈 리스트를 생성한 뒤, heapq 모듈을 통해 원소를 추가하거나 삭제한 리스트가 최소 힙에 해당한다.
heap = []
원소 추가
heapq.heappush 함수를 이용하여 힙에 원소를 추가한다. 첫 번째 매개변수는 원소를 추가할 리스트이고, 두 번째 매개변수는 추가한 원소이다.
heapq.heappush(heap, val)
# 예시
heapq.heappush(heap, 4)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
원소 삭제
heaqp.heappop 함수를 이용하여 힙에서 원소를 삭제한다. 원소를 삭제할 리스트를 매개변수로 하여 실행하며, 리스트에서 가장 작은 원소를 삭제한 후 값을 반환한다.
heapq.heappop(heap)
# 예시
print(heapq.heappop(heap))
최솟값 읽기
원소를 삭제하지 않고, 최소값을 읽고 싶다면 인덱스를 통해 리스트에 접근하면 된다.
heap[0]
# 예시
print(heap[0])
리스트를 힙으로 변환
heapq.heapify 함수를 사용하여 리스트를 힙으로 변환할 수 있다. 매개변수로 전달한 리스트는 힙 구조에 맞게 재배치된다.
heapq.heapify(heap)
# 예시
heap = [4, 1, 7, 3, 8, 5]
heapq.heapify(heap)
print(heap)
최대 힙
heapq는 최소 힙으로만 동작하기 때문에, 최대 힙처럼 사용하기 위해서는 약간의 요령이 필요하다.
한 가지는 튜플을 활용하는 방법이고, 다른 한 가지는 (원소가 숫자인 경우) 음수로 변환하여 힙에 저장하는 것이다.
음수 활용 저장
heap = []
heapq.heappush(heap, -1 * 4)
heapq.heappush(heap, -1 * 1)
heapq.heappush(heap, -1 * 7)
print(*map(abs, heap))