[파이썬] heapq 사용법 (우선순위 큐, 힙)

파이썬에서는 우선순위가 높은 데이터를 먼저 꺼낼 수 있는 heapq라는 모듈을 제공합니다. (값이 낮은 순으로 뽑는 최소 힙으로 구현되어있으며 최대 힙은 따로 만들어져있지 않고 최소 힙을 응용해서 사용할 수 있습니다.)

다른 언어에서는 우선순위 큐를 제공하기도 하는데 파이썬에는 없는 것으로 알고 있습니다.

 

힙은 이진트리로 만들어져 있어 삽입, 삭제를 하는데 O(log n)이 걸린다는 것을 알 수 있습니다.

 

1. 힙 생성 및 삽입하기

import heapq

h = []
heapq.heappush(h, (3, "third"))
heapq.heappush(h, (2, "second"))
heapq.heappush(h, (4, "fourth"))
heapq.heappush(h, (1, "first"))

print(h)

#[(1, 'first'), (2, 'second'), (4, 'fourth'), (3, 'third')]

먼저 모듈을 import해주고 삽입할 배열을 만듭니다.

그리고 heappush함수를 사용하는데요 첫 번째 인자로는 삽입할 리스트, 두 번째 인자로는 튜플을 사용할 경우 튜플의 첫 번째 항목을 기준으로 우선순위를 가지게 됩니다. 그냥 정수 값으로 할 수도 있습니다.

 

2. pop하기

import heapq

h = []
heapq.heappush(h, (3, "third"))
heapq.heappush(h, (2, "second"))
heapq.heappush(h, (4, "fourth"))
heapq.heappush(h, (1, "first"))

print(heapq.heappop(h))
print(heapq.heappop(h))
print(heapq.heappop(h))
print(heapq.heappop(h))

#(1, 'first')
#(2, 'second')
#(3, 'third')
#(4, 'fourth')

pop을 하게 되면 우선순위가 가장 높은 항목을 리턴하면서 삭제하게 됩니다.(우선순위가 높다 = 값이 작다)

heappop함수의 인자로는 리스트를 전달해주면 됩니다.

 

3. 주어진 배열을 heap으로 바꾸기

import heapq

h = [ [4, "fourth"], [2, "second"], [3, "third"], [1, "first"]]
heapq.heapify(h)
print(h)

#[[1, 'first'], [2, 'second'], [3, 'third'], [4, 'fourth']]

일일이 n개의 데이터를 삽입한다고 했을 때 각 삽입마다 O(log n)의 시간이 걸리므로 총 O(nlog n)의 시간이 걸림을 알 수 있습니다. 하지만 heapify를 이용할 경우 O(n)의 시간으로 배열을 heap으로 바꿔줄 수 있습니다.

 

4. 최대 힙

앞서 말했다시피 힙에서 두 번째 인자의 첫 번째 항목을 기준으로 우선순위를 정하게 됩니다. 1->2->3.... 순으로 우선순위가 결정되는데 3->2->1로 바꾸고 싶으면 값에 -를 붙이면 됩니다.

import heapq

h = [ [-4, 4, "fourth"], [-2, 2, "second"], [-3, 3, "third"], [-1, 1, "first"]]
heapq.heapify(h)
print(heapq.heappop(h))
print(heapq.heappop(h))
print(heapq.heappop(h))
print(heapq.heappop(h))

#[-4, 4, 'fourth']
#[-3, 3, 'third']
#[-2, 2, 'second']
#[-1, 1, 'first']