Python의 heapq 모듈은 힙(Heap) 자료구조를 구현하는 모듈로, 우선순위 큐(Priority Queue)를 효율적으로 구현할 수 있게 해줍니다.
heapq 모듈 특징
- 최소 힙(Min Heap): 가장 작은 값이 루트에 위치
- 이진 힙(Binary Heap): 완전 이진 트리 구조 (트리 자료 구조)
- O(log n): 삽입, 삭제 시간복잡도
- O(1): 최솟값 조회 시간복잡도
- 리스트 기반: Python 리스트를 힙으로 변환하여 사용
import 방법
import heapq
# 또는
from heapq import heappush, heappop, heapify, heappushpop, heapreplace
주요 함수 정리
| 함수 | 설명 | 시간복잡도 |
|---|
heapify(heap) | 리스트를 힙으로 변환 (in-place) | O(n) |
heappush(heap, item) | 힙에 요소 추가 | O(log n) |
heappop(heap) | 힙에서 최솟값 제거하고 반환 | O(log n) |
heappushpop(heap, item) | 힙에 item 추가 후 최솟값 반환 (heappush + heappop) | O(log n) |
heapreplace(heap, item) | 힙에서 최솟값 제거 후 item 추가 (heappop + heappush) | O(log n) |
nlargest(n, iterable, key=None) | iterable에서 가장 큰 n개 요소 반환 | O(n log k) |
nsmallest(n, iterable, key=None) | iterable에서 가장 작은 n개 요소 반환 | O(n log k) |
기본 사용법
import heapq
# 빈 힙 생성
heap = []
# 요소 추가
heapq.heappush(heap, 4)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3)
print(heap) # [1, 3, 7, 4] (힙 구조)
# 최솟값 제거 및 반환
print(heapq.heappop(heap)) # 1
print(heap) # [3, 4, 7]
# 기존 리스트를 힙으로 변환
numbers = [5, 2, 8, 1, 9]
heapq.heapify(numbers)
print(numbers) # [1, 2, 8, 5, 9] (힙 구조)
heappush와 heappop
import heapq
heap = []
# heappush: 요소 추가
heapq.heappush(heap, 10)
heapq.heappush(heap, 5)
heapq.heappush(heap, 15)
heapq.heappush(heap, 3)
print(heap) # [3, 5, 15, 10]
# heappop: 최솟값 제거 및 반환
print(heapq.heappop(heap)) # 3
print(heap) # [5, 10, 15]
print(heapq.heappop(heap)) # 5
print(heap) # [10, 15]
heapify
import heapq
# 기존 리스트를 힙으로 변환
numbers = [9, 2, 7, 1, 5, 8, 3]
heapq.heapify(numbers)
print(numbers) # [1, 2, 3, 9, 5, 8, 7]
# 힙 구조 확인
print(heapq.heappop(numbers)) # 1
print(heapq.heappop(numbers)) # 2
print(heapq.heappop(numbers)) # 3
heappushpop과 heapreplace
import heapq
heap = [3, 5, 7, 9]
# heappushpop: 추가 후 최솟값 반환
result = heapq.heappushpop(heap, 2)
print(result) # 2 (추가된 값이 최솟값이었음)
print(heap) # [3, 5, 7, 9]
result = heapq.heappushpop(heap, 10)
print(result) # 3 (기존 최솟값 반환)
print(heap) # [5, 9, 7, 10]
# heapreplace: 최솟값 제거 후 추가
result = heapq.heapreplace(heap, 1)
print(result) # 5 (제거된 최솟값)
print(heap) # [1, 9, 7, 10]
nlargest와 nsmallest
import heapq
numbers = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
# 가장 큰 3개 요소
largest = heapq.nlargest(3, numbers)
print(largest) # [9, 6, 5]
# 가장 작은 3개 요소
smallest = heapq.nsmallest(3, numbers)
print(smallest) # [1, 1, 2]
# key 함수 사용 (문자열 길이 기준)
words = ['apple', 'banana', 'cherry', 'date', 'elderberry']
longest = heapq.nlargest(2, words, key=len)
print(longest) # ['elderberry', 'banana']
튜플을 이용한 우선순위 큐
import heapq
# (우선순위, 데이터) 형태로 저장
pq = []
heapq.heappush(pq, (3, 'task3'))
heapq.heappush(pq, (1, 'task1'))
heapq.heappush(pq, (2, 'task2'))
# 우선순위가 낮은 순서대로 처리
while pq:
priority, task = heapq.heappop(pq)
print(f"Processing {task} with priority {priority}")
# 출력:
# Processing task1 with priority 1
# Processing task2 with priority 2
# Processing task3 with priority 3
최대 힙 구현
import heapq
# 최대 힙: 음수로 변환하여 저장
max_heap = []
numbers = [3, 1, 4, 1, 5, 9, 2, 6]
# 음수로 변환하여 최소 힙에 저장
for num in numbers:
heapq.heappush(max_heap, -num)
# 최댓값 추출 (음수로 저장했으므로 다시 양수로 변환)
print(-heapq.heappop(max_heap)) # 9
print(-heapq.heappop(max_heap)) # 6
print(-heapq.heappop(max_heap)) # 5
Top K 문제
import heapq
def find_top_k(nums, k):
"""가장 큰 k개 요소 찾기"""
return heapq.nlargest(k, nums)
def find_bottom_k(nums, k):
"""가장 작은 k개 요소 찾기"""
return heapq.nsmallest(k, nums)
# 예시
numbers = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
print(find_top_k(numbers, 3)) # [9, 6, 5]
print(find_bottom_k(numbers, 3)) # [1, 1, 2]
주요 팁
- 최소 힙만 지원: 최대 힙이 필요하면 음수로 변환하여 사용
- 리스트 기반: 기존 리스트를 heapify()로 힙으로 변환 가능
- 안정성: 같은 우선순위의 요소들은 삽입 순서대로 처리되지 않을 수 있음
- 메모리 효율성: 리스트를 직접 힙으로 사용하므로 메모리 효율적
- 정렬과의 차이: 전체 정렬이 아닌 최솟값만 빠르게 접근
주의사항
- 힙 구조 유지: heappush, heappop 외의 방법으로 리스트 수정 시 힙 구조가 깨질 수 있음
- 중복 요소: 같은 값이 여러 개 있을 수 있음
- 정렬되지 않음: 힙은 완전히 정렬된 상태가 아님 (부모-자식 관계만 보장)
- 인덱스 접근: heap[0]으로 최솟값만 안전하게 접근 가능
성능 비교
| 연산 | heapq | 정렬된 리스트 | 일반 리스트 |
|---|
| 최솟값 접근 | O(1) | O(1) | O(n) |
| 삽입 | O(log n) | O(n) | O(1) |
| 삭제 | O(log n) | O(n) | O(n) |
| 전체 정렬 | O(n log n) | O(n log n) | O(n log n) |
언제 사용할까?
- 우선순위 큐: 작업 스케줄링, 이벤트 처리
- Top K 문제: 가장 큰/작은 k개 요소 찾기
- 중간값 찾기: 실시간 데이터 스트림에서 중간값 계산
- 최단 경로: 다익스트라, 프림 알고리즘
- 힙 정렬: O(n log n) 정렬 알고리즘 구현