[파이썬] heapq

파이썬의 heap queue 알고리즘

완전 이진 트리는 heaq을 이용해서 구현하는데, heaq은 배열(리스트)을 이용해서 구현한다.

따라서, heap을 생성할 때에는 []로 초기화한 list를 이용한다.

heapq 사용

파이썬에는 heapq라는 내장 모듈이 존재한다. 이 heapq를 이용해서 우선순위 큐를 구현할 수 있다. PriorityQueue라는 모듈도 존재하는데, heapq이 성능 측면에서 훨씬 좋다.

파이썬에서 heapq를 사용하기 위해서는 아래와 같이 import해주어야 한다.

import heapq

heapq 생성

힙은 배열을 이용해서 구현하기 때문에, 최소힙을 만들거라면 배열을 초기화한다.

import heapq

queue = []

heapq push

heapq.heappush(heap, item)

힙에 push하기 위해서는 heapq 모듈에 있는 heappush라는 함수를 사용해야 한다.

import heapq

queue = []

heapq.heappush(queue, 1)
heapq.heappush(queue, 2)
heapq.heappush(queue, -3)

heapq pop

heapq.heappop(heap)

힙에서 가장 작은 원소를 pop하고 싶다면 heapq 모듈에 있는 heappop이라는 함수를 사용해야 한다.

import heapq

queue = []

heapq.heappush(queue, 1)
heapq.heappush(queue, 2)
heapq.heappush(queue, -3)

element = heapq.heappop(queue) # -3

heapq push하고 바로 pop

heapq.heappushpop(heap, item)

힙에서 push하자마자 pop할 수 있는 함수도 있다. heappush + heappop하는 것보다 훨씬 효율적이다.

import heapq

queue = []

heapq.heappush(queue, 1)
heapq.heappush(queue, 2)
heapq.heappush(queue, -3)

element = heapq.heappushpop(queue, -4) # -4 

heapq []를 바로 heap으로

heapq.heapify(list)

list x를 heap으로 바꿔준다. 이때 linear-time 이 소요된다.

import heapq

queue = [-4, 1, 2, -3]

heapq.heapify(queue)

알고리즘 문제 풀 때 팁

힙 안에 넣은 모든 요소를 밖으로 빼낼 때, while heap이라고 하면 쉽게 꺼낼 수 있다.

import heapq

queue = []

heapq.heappush(queue, 1)
heapq.heappush(queue, 2)
heapq.heappush(queue, -3)
heapq.heappush(queue, -4)

while queue:
	element = heapq.heappop(queue)
    print(element)

# 결과 (가장 작은 값부터 pop되는 것을 볼 수 있다.)
# -4
# -3
# 1
# 2

참고 사이트: https://docs.python.org/3/library/heapq.html