[python에서 priority queue 자료구조 사용하기 / heapq]

코딩 문제 하나를 풀었는데, 문제 풀이의 핵심이 '가지고 있는 값들 중에 가장 큰 값 두 개를 반복적으로 꺼내는 것'이였다. 이 문제를 풀려면 priority queue 자료구조를 사용하는 것이 효과적이라고 판단했다. 파이썬에서는 priority queue를 heapq로 편하게 사용할 수 있다. 그래서 오늘은 이 heapq에 대해 공부해봤다.

---

1. 먼저, priority queue를 사용해서 이 문제를 푸는 데 걸리는 시간은 O(nlogn)이다.

2. heapq를 사용하려면 먼저 import heapq를 해야 한다.

3. heapq 라이브러리의 주요 메서드
- heapq.heapify(list): 리스트를 heap으로 변환한다.
- heapq.heappop(heap): heap에서 가장 작은 원소를 pop, 리턴. 비어있는 경우는 IndexError.
- heapq.heappush(heap, item): item을 heap에 추가

4. heap의 0번 인덱스는 최솟값이다.

 

5. heapq는 기본적으로 최소 힙(min-heap)이다.

- 따라서 heappop()을 하면 가장 작은 값이 먼저 나온다.

- 만약 최대 힙(max-heap)을 만들고 싶다면 음수 값을 저장하는 방식을 써야 한다.

import heapq

nums = [1, 3, 5, 7, 9]
max_heap = [-num for num in nums]  # 모든 값에 -를 붙임
heapq.heapify(max_heap)

print(-heapq.heappop(max_heap))  # 9 (최대값)

 

6. heapq.nlargest(n, iterable)과 heapq.nsmallest(n, iterable)을 활용하면 더 쉽게 최대/최소 n개를 뽑을 수 있다.

import heapq

nums = [1, 3, 5, 7, 9, 2, 4, 6, 8]
print(heapq.nalrgest(2, nums))  # [9, 8]
print(heapq.nsmallest(2, nums))  # [1, 2]

- '가장 큰 값 두 개를 반복적으로 꺼내는 문제'에선 nlargest(2, heap)도 한 방법이 될 수 있다.

 

7. 우선순위 큐로 활용할 때 튜플을 사용할 수 있다.

- 예를 들어 (우선순위, 값) 형식으로 저장하면 자동으로 우선순위가 낮은 것부터 처리된다.

import heapq

tasks = [(2, "write code"), (1, "study"), (3, "exercise")]
heapq.heapify(tasks)

print(heapq.heappop(tasks))  # (1, "study") -> 우선순위가 가장 낮은 값이 먼저 나옴