[Python] heapq 라이브러리

힙큐(heapq)는 파이썬 표준라이브러리가 제공하는 모듈로 Heap 자료구조를 지원한다. 
완전 이진트리 형태를 하고 있으며, 최솟값과 최댓값을 빠르게 찾을 수 있다 . 
파이썬은 기본적으로 최소힙(min-heap)을 제공한다. (부모 노드의 값은 자식 노드의 값보다 작거나 같다.)

주요 메소드

heapq.heappush(heap, item)

힙의 조건을 유지하면서 item을 heap에 push 해준다. ( O(logn) )

import heapq

heap = []

for i in range(10):
    heapq.heappush(heap, i)
print(heap)

## [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

heapq.heappop(heap)

힙에서 가장 작은 요소를 제거하고 반환함.  ( O(logn) )

힙이 비어있으면 IndexError을 발생시킴

import heapq

heap = [3,4,5,6]

print(heapq.heappop(heap))
### 3
print(heap)
### [4, 6, 5]

heapq.heapify(list)

리스트를 힙으로 변환시킴 ( O(n) )

리스트를 최소 힙 속성을 가지도록 정렬함. 

import heapq

list = [3,6,34,3,5,7,3,6,24,143]

heapq.heapify(list)

print(list)
##[3, 3, 3, 6, 5, 7, 34, 6, 24, 143]

heapq.nlargest ( n, heap, key=None )

heap에서 가장 큰 n 개의 요소를 반환함. ( O(nlogn) )

key를 통해 정렬기준을 정할 수 있다. 

import heapq
heap = []
for i in range(10):
    heapq.heappush(heap, i)

print(heapq.nlargest(3,heap))
## [9, 8, 7]

 

heapq.nsmallest(n, heap, key=None)

heap에서 가장 작은 n개의 요소를 반환( O(nlogn) )

key를 통해 기준 설정 가능

import heapq
heap = []
for i in range(10):
    heapq.heappush(heap, i)

print(heapq.nsmallest(3,heap))
#[0, 1, 2]

활용 예시

1. 최대 힙 구현

import heapq

heap = []
number = [4,3,2,3,1]

for num in number:
	heapq.heappush(heap, -num) #[-4, -3, -3, -2, -1]

for _ in range(len(heap)):
	print(-heapq.heappop(heap), end = ", ")

## 4 3 3 2 1

파이썬에서는 최대 힙을 제공하지 않는다. 

힙에 부호를 변경하여 값을 넣어주면 간접적으로 최대 힙을 구현할 수 있다. 

2. 중앙 값 찾기

import heapq

class MedianFinder:
    def __init__(self):
        self.left_heap = []  # 최대 힙
        self.right_heap = []  # 최소 힙

    def add_num(self, num):
        # 최대 힙에 추가 
        heapq.heappush(self.left_heap, -num)
        
        # 최대 힙의 최대값을 최소 힙에 이동
        heapq.heappush(self.right_heap, -heapq.heappop(self.left_heap))
        
        # 최소 힙의 크기가 더 커지면 최소 힙의 최소값을 최대 힙으로 이동
        if len(self.right_heap) > len(self.left_heap):
            heapq.heappush(self.left_heap, -heapq.heappop(self.right_heap))

    def find_median(self):
        if len(self.left_heap) > len(self.right_heap):
            return -self.left_heap[0]
        else:
            return min(-self.left_heap[0], self.right_heap[0])


medianFinder = MedianFinder()
numbers = [4,1,3,4,5,6,4,3,3]

for num in numbers:
    medianFinder.add_num(num)
    print(f"Inserted {num}, current median is: {medianFinder.find_median()}")
    
"""
Inserted 4, current median is: 4
Inserted 1, current median is: 1
Inserted 3, current median is: 3
Inserted 4, current median is: 3
Inserted 5, current median is: 4
Inserted 6, current median is: 4
Inserted 4, current median is: 4
Inserted 3, current median is: 4
Inserted 3, current median is: 4
"""

최대 힙과 최소 힙을 활용하여 중앙 값 찾는 알고리즘을 구현할 수 있다.