heapq

파이썬에서는 heapq 모듈을 사용해서 최소 힙과 최대 힙을 구현할 수 있다.

heapq 내부 메소드

heapq 모듈의 내부 메소드들 중 주로 쓰이는 메소드는 heappush와 heappop이 있다.

heappush(heap, item)

힙 불변성을 유지하면서, item 값을 heap으로 push 해주는 메소드이다.

heappop(heap)

힙 불변성을 유지하면서, heap에서 가장 작은 항목을 pop하고 반환해주는 메소드이다. heappop 메소드를 사용할 때 주의해야할 점은, 힙이 비어 있을 때 heappop 메소드를 호출하면 IndexError가 발생한다는 것이다.

최소 힙

다음과 같이 최소 힙을 구현할 수 있다.

heap 역할을 할 빈 리스트를 선언해주고, heappush 메소드의 매개변수로 해당 리스트와 추가하고자 하는 값을 넣어주면 리스트에 힙 정렬이 되게 원하는 값이 삽입된다.

import heapq

heap = []

# 아래 for문을 실행하고 나면 heap은 [1,2,3,5,4]로 힙 정렬이 되게 된다.
for i in range(1,6):
	heapq.heappush(heap, i)

# 작은 숫자 순서대로 1,2,3,4,5가 출력된다.
for i in range(5):
	print(heapq.heappop(heap))

최대 힙

heapq에서는 최대 힙을 제공하지 않는다. 따라서 다음과 같이 부호를 변경하는 방법을 사용해서 최대 힙을 구현한다. 부호를 바꿔서 최소 힙에 넣어준 뒤에 최솟값부터 pop을 해줄 때 다시 부호를 바꿔주면 최대 힙과 동일하다.

import heapq

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

# 아래 for문을 실행시키고 나면 heap은 [-5,-4,-3,-1,-2]가 된다.
for value in values:
	heapq.heappush(heap, -value)

# 아래 for문을 실행시키면 5,4,3,2,1이 출력된다. 즉, 큰 숫자부터 출력이 된다.
for i in range(5):
	print(-heapq.heappop(heap))