파이썬에서 우선순위 큐 사용하기

때때로 프로그래밍 문제를 풀다보면 우선순위 큐를 활용해야 하는 경우가 종종 있습니다.

다만 우선순위 큐는 일반적인 큐나 배열이 아닌 힙을 기반으로 구현되었기 때문에 이를 직접 구현해서 사용하기에는 시간이 조금 걸릴수도 있는데요, 다행히 파이썬에서는 우선순위 큐를 내장 라이브러리로 제공합니다.

⏰ 들어가기 전에

우선순위 큐를 구현한 PriorityQueue 라는 라이브러리도 존재하지만, heapq 보다 훨씬 느리게 동작하므로 코딩 테스트용으로는 적합하지 않습니다.

따라서 이번 글에서는 heapq 를 사용한 방법을 소개하도록 하겠습니다.

import heapq

# 일반적인 리스트를 선언합니다.
pq = []

우선순위 큐는 원소를 삽입할 때는 일반적인 큐와 유사하지만, 원소를 큐에서 뺄 때는 언제나 작은 값을 먼저 내보낸다는 특징이 있습니다.

import heapq

pq = []

heapq.heappush(pq, 1)
heapq.heappush(pq, 5)
heapq.heappush(pq, 2)
heapq.heappush(pq, 4)
heapq.heappush(pq, 3)

print(heapq.heappop(pq)) # 1
print(heapq.heappop(pq)) # 2
print(heapq.heappop(pq)) # 3
print(heapq.heappop(pq)) # 4
print(heapq.heappop(pq)) # 5

또, 둘 이상의 비교대상을 튜플 형태의 인자로 받으면 튜플의 원소를 차례로 비교해가며 가장 작은 값을 판별합니다.

import heapq

pq = []

heapq.heappush(pq, (1, 2, 4))
heapq.heappush(pq, (2, 1, 1))
heapq.heappush(pq, (1, 3, 4))
heapq.heappush(pq, (1, 3, 3))
heapq.heappush(pq, (1, 2, 3))

print(heapq.heappop(pq)) # (1, 2, 3)
print(heapq.heappop(pq)) # (1, 2, 4)
print(heapq.heappop(pq)) # (1, 3, 3)
print(heapq.heappop(pq)) # (1, 3, 4)
print(heapq.heappop(pq)) # (2, 1, 1)

우선순위 큐의 사용법을 연습해보고 싶다면, 아래 3문제를 한번 풀어보시는걸 추천드립니다 :)

1. 14593번 - 2017 아주대학교 프로그래밍 경시대회 (Large)
2. 11279번 - 최대 힙
3. 1715번 - 카드 정렬하기

부록 : PriorityQueue와 heapq의 실행속도 비교

# heapq와 PriorityQueue의 속도 차이를 비교하기 위한 코드입니다.

from queue import PriorityQueue
from time import time
import heapq
import random

nums = [random.random() for _ in range(100000)]

pq_priority_queue = PriorityQueue()
pq_heapq = []

start = time()
for i in range(len(nums)):
    pq_priority_queue.put(nums[i])

for i in range(len(nums)):
    pq_priority_queue.get()
end = time()
print(f'PriorityQueue 동작 시간 : {end - start:.6f}')

start = time()
for i in range(len(nums)):
    heapq.heappush(pq_heapq, nums[i])

for i in range(len(nums)):
    heapq.heappop(pq_heapq)
end = time()

print(f'heapq 동작 시간 : {end - start:.6f}')

PriorityQueue가 더 편하고 직관적이기는 하지만, PriorityQueue를 쓰면 시간 초과가 나는 문제가 heapq를 사용하면 통과되는 경우가 꽤 있습니다.

 

따라서 코딩 테스트가 주 목적이라면 반드시 heapq를 사용할 것을 권해드립니다. 😄