[파이썬] 우선순위 큐(priority queue)를 위한 heapq 모듈 사용법

5 min readFeb 9, 2019

우선순위 큐는 우선 순위가 가장 높은 자료(data)를 가장 먼저 꺼낼 수 있는 자료 구조이다. 배열을 사용하면 직접 구현하기는 어렵지 않지만, 파이썬에서는 heapq라는 내장(built-in) 모듈로 제공이 되기 때문에, 추가적인 연산이 필요 없다면 내장 모듈을 사용하는게 좋다.

*이 글에서 힙은 이진힙을 의미하고, 우선순의 큐는 개별의 값을 가지는 힙이므로 두 단어를 번갈아 사용하겠다.

우선 순위 큐의 생성 및 원소 삽입

heapq.heappush를 사용해 우선 순위 큐 또는 힙을 쉽게 생성할 수 있다. 첫번째 인자는 heap 자체인 list이고, 두번째 인자는 튜플인데 튜플의 첫번째 요소는 우선순위 값, 두번째 요소는 데이터를 넣어주면 된다. 함수의 두번째 인자로 튜플이 아닌 일반 값을 넣어주면, 값을 기준으로 heap을 만들어준다. 파이썬이 제공하는 힙은 최소힙(min-heap)이므로 주의하자. 삽입 별 시간 복잡도는 O(log n) 이다.

다음 예제를 보자.

import heapqh = []
heapq.heappush(h, (3, "Go to home"))
heapq.heappush(h, (10, "Do not study"))
heapq.heappush(h, (1, "Enjoy!"))
heapq.heappush(h, (4, "Eat!"))
heapq.heappush(h, (7, "Pray!"))print(h)

결과는 배열로 구현된 힙을 보여준다.

[(1, 'Enjoy!'), (4, 'Eat!'), (3, 'Go to home'), (10, 'Do not study'), (7, 'Pray!')]

2. 원소 꺼내기

원소 꺼내기는 heapq.heappop 으로 구현하면 된다. 각 꺼내기 별 시간 복잡도는 O(log n)이다.

import heapqh = []
heapq.heappush(h, (3, "Go to home"))
heapq.heappush(h, (10, "Do not study"))
heapq.heappush(h, (1, "Enjoy!"))
heapq.heappush(h, (4, "Eat!"))
heapq.heappush(h, (7, "Pray!"))first = heapq.heappop(h)
second = heapq.heappop(h)
third = heapq.heappop(h)print("first:", first)
print("second:", second)
print("third:", third)

우선순위 순서대로 나온 결과를 확인할 수 있다.

first: (1, 'Enjoy!')
second: (3, 'Go to home')
third: (4, 'Eat!')

3. 배열로 부터 힙 정렬 만들기

1번의 방법으로 배열을 힙으로 만들면 시간 복잡도는 O(n log n)이다. 그러나 배열로 부터 힙을 만드는 최적의 알고리즘의 시간 복잡도는 O(n) 으로 알려져 있다. 파이썬에서는 O(n)의 시간으로 배열을 힙으로 만들 수 있는 heapq.heapfy 함수를 제공한다. 주의할 점은 배열 자체가 힙으로 바뀐다는 점이다.

import heapqh = [(3, "Go to home"), (10, "Do not study"), (1, "Enjoy!"), (4, "Eat!"), (7, "Pray!")]
heapq.heapify(h)
print(h)

결과는 1번과 동일하다.

[(1, 'Enjoy!'), (4, 'Eat!'), (3, 'Go to home'), (10, 'Do not study'), (7, 'Pray!')]

4. 최대 힙 (max-heap)

눈치 챈 사람도 있겠지만, heapq 는 기본적으로 최소 힙 밖에 지원을 안한다. (왜 제공을 안해주는지 모르겠다.) 최대 힙을 쓸 수 있는 방법은 없을까? heapq 에서 공식적으로 지원해주는 기능은 없다. 인터넷에서 검색을 해보면, heapq._heapfy_max 나 heapq._heappop_max 를 사용해서 하는 방법도 있지만, push를 지원해주지 않기 때문에, 반쪽짜리 기능이다. 유일한 방법은 키 값을 변환해서 넣는 수 밖에 없다. 모든 요소의 키를 바꾸는데 드는 시간은 O(n) 이고, 모든 힙 요소를 얻는데 드는 시간 복잡도는 O(n log n) 이므로, 요소를 모두 돌면서 키를 바꾼다 해도 전체 시간 복잡도에 문제가 되지 않는다. 앞의 예에서는 키 값이 정수 이므로 -를 붙여주면 된다.

import heapqh = [(3, "Go to home"), (10, "Do not study"), (1, "Enjoy!"), (4, "Eat!"), (7, "Pray!")]
max_h = [(-ele[0], ele[1]) for ele in h]
heapq.heapify(max_h)
print(max_h)

그러나 이 방법은 개별적으로 push를 사용할 때, 매번 데이터를 변경해 줘야 하는 번거로움이 있다. 이러한 번거로움을 피하고 싶으면, 결국 heapq 모듈의 기능을 래핑해 max 힙을 만드는 클래스를 정의하는 수 밖에 없다. __lt__ 내장 메소드를 활용하면, 비교는 반대로 할 수 있으므로, 어렵지 않게 구현할 수 있다.

주의할 점은 위 두가지 방법 다 원본 데이터를 유지할 수 있는 방법을 제공해 주어야 한다는 점이다. 이를 위해서 데이터의 래핑와 언래핑을 하도록 구현을 해야 한다.

다음 코드는 __lt__ 내장 메서드를 활용해 max_heapq 모듈을 간단히 구현해본 것이다. 요소를 push 하거나 pop 할때, 요소를 래핑하고 언래핑해 주고 있음에 유의하자.

두번째 방법은 힙을 직접 구현하는 것이다. 배열을 사용한 이진 힙 구현은 쉽기 때문이다. 내장 모듈 보다 실제 속도는 느리겠지만, 시간 복잡도가 변하지는 않는다.

[파이썬] 우선순위 큐(priority queue)를 위한 heapq 모듈 사용법

Written by Yunhong Min