[heap] 파이썬으로 힙 구현을 위한 heapq 라이브러리

Python heapq 모듈 정리

heapq 모듈은 Python에서 힙(heap) 자료구조를 제공하며, 효율적으로 최소값을 추출하거나 우선순위 큐를 구현할 때 사용됩니다. 기본적으로 최소 힙(min-heap)을 지원합니다.

1. 힙의 특징

• 최소 힙: 루트(0번 인덱스) 노드가 항상 최솟값을 유지합니다.
• 최대 힙은 기본적으로 제공되지 않으며, 음수 값을 이용해 구현할 수 있습니다.

2. 주요 함수

1) heapq.heappush(heap, item)

힙에 새로운 요소를 추가합니다. 시간 복잡도는 O(log n)입니다.

import heapq

heap = []
heapq.heappush(heap, 10)
heapq.heappush(heap, 1)
heapq.heappush(heap, 5)
print(heap)  # [1, 10, 5]

2) heapq.heappop(heap)

힙에서 최솟값을 제거하고 반환합니다. 시간 복잡도는 O(log n)입니다.

print(heapq.heappop(heap))  # 1
print(heap)  # [5, 10]

3) heapq.heappushpop(heap, item)

새로운 요소를 추가한 뒤, 힙에서 최솟값을 바로 제거하고 반환합니다. 더 효율적인 방식으로 O(log n) 시간 복잡도를 가집니다.

heap = [10, 20, 30]
heapq.heapify(heap)
print(heapq.heappushpop(heap, 5))  # 5
print(heap)  # [10, 20, 30]

4) heapq.heapreplace(heap, item)

힙의 루트를 제거한 뒤 새로운 요소로 교체합니다. 시간 복잡도는 O(log n)입니다.

print(heapq.heapreplace(heap, 25))  # 10
print(heap)  # [20, 25, 30]

5) heapq.heapify(iterable)

리스트를 힙 구조로 변환합니다. 시간 복잡도는 O(n)입니다.

nums = [3, 2, 5, 4, 1]
heapq.heapify(nums)
print(nums)  # [1, 2, 5, 4, 3]

6) 최대 힙 구현

음수 값을 이용하거나 커스텀 정렬을 통해 최대 힙을 구현할 수 있습니다.

max_heap = []
heapq.heappush(max_heap, -10)
heapq.heappush(max_heap, -1)
heapq.heappush(max_heap, -5)
print([-x for x in max_heap])  # [10, 1, 5]
print(-heapq.heappop(max_heap))  # 10

3. 활용 예시

1) K개의 가장 작은 요소 추출

nums = [7, 2, 4, 10, 3, 1]
heapq.heapify(nums)
k = 3
smallest_k = [heapq.heappop(nums) for _ in range(k)]
print(smallest_k)  # [1, 2, 3]

2) 우선순위 큐

tasks = [(5, 'low priority'), (1, 'high priority'), (3, 'medium priority')]
heapq.heapify(tasks)
while tasks:
    priority, task = heapq.heappop(tasks)
    print(f"{task} with priority {priority}")
# 출력:
# high priority with priority 1
# medium priority with priority 3
# low priority with priority 5

3) K개의 가장 큰 요소 추출

heapq.nlargest를 사용하여 손쉽게 구현 가능합니다.

nums = [7, 2, 4, 10, 3, 1]
print(heapq.nlargest(3, nums))  # [10, 7, 4]