파이썬에서 우선순위 큐(Priority Queue)를 구현하는 방법

Programmer 문제를 풀던 중 Priority Queue가 필요한 경우가 생겼다.

파이썬에는 여러 라이브러리가 있다보니 대중적으로 알고리즘 문제 풀이에 많이 쓰이는

Priority Queue 구현 방법을 복기하고자 한다.

---

heapq 모듈

파이썬에서는 heapq 모듈을 사용해 간단하고 효율적으로 리스트 자료형의 우선순위 큐를 구현할 수가 있다.

 

heapq가 단순히 우선순위 큐만을 제공하기 위해 나온 것은 아니고, 최소 힙(Min-Heap) 자료구조를 제공하는 모듈이다.

이를 통해 우선순위 큐를 구현할 수 있는 이유는 최소 힙은 완전 이진 트리로, 부모 노드가 항상 자식 노드보다 작거나 같은 값을 가지는 특성을 갖고 있기 때문이다.

 

이를 통해서 항상 최솟값을 O(1)의 시작 복잡도로 얻을 수 있고, 트리 노드를 구성해야 하는만큼 삽입 및 삭제 연산에는 O(log n)의 시간 복잡도를 가진다.

 

파이썬에서는 따로 객체 없이 리스트 형식으로 heapq를 구현하여 우선순위 큐를 쉽게 만들 수가 있다.

heapq를 사용한 우선순위 큐 구현

그럼 이제 우선순위 큐를 구현하기 위해 heapq가 제공하는 기본 함수들을 활용해보자.

먼저 heapq를 사용하기 위해서는 heapq 모듈을 import 해주어야 한다.

import heapq

 

생성 및 삽입

힙을 생성하는 것은 어렵지 않다. 그냥 기본으로 사용되는 파이썬의 리스트를 생성하면 된다.

이후 힙에 들어갈 값을 heapq.heappush(heap_list, element) 함수를 통해 삽입해준다.

 

그냥 리스트에 요소를 추가하듯이 pq.append(element)로 삽입하면 해당 요소는 힙 노드 요소로 들어가지 않기 때문에 위 함수를 사용해야 한다.

pq = [] #힙(우선순위 큐)으로 사용하기 위한 빈 리스트 생성

#우선순위 큐에 요소 추가
heapq.heappush(pq, 5)
heapq.heappush(pq, 3)
heapq.heappush(pq, 6)

 

삭제

가장 작은 요소를 반환하며 삭제하는 이른 바 pop을 하기 위해서는 heapq.heappop(heap_list) 함수를 사용한다. 

이 함수는 힙에서 가장 작은 요소를 반환하며 해당 요소를 삭제한다.

print(heapq.heappop(pq)) #3
print(heapq.heappop(pq)) #5
print(heapq.heappop(pq)) #6

 

리스트를 힙으로

값을 직접 삽입하지 않고, 이미 존재하는 리스트의 요소들을 힙 요소로 바꾸고 싶은 경우도 있을 것이다.

그럴 때엔 heapq.heapify(heap_list) 함수를 사용해 해당 리스트 및 요소들을 힙으로 변경해 준다.

origin_list = [6, 2, 3]
heapq.hepaify(origin_list) #해당 리스트의 요소들을 힙 노드들로 변경

print(heapq.heappop(pq)) #2
print(heapq.heappop(pq)) #3
print(heapq.heappop(pq)) #6