[Python] heapq 라이브러리

파이썬의 heapq 라이브러리에 대해서 알아보자!

 

1. 소개

heap 자료구조를 구현하는 내장 모듈이고 최소 힙을 기본으로 제공한다.

따라서 가장 작은 원소가 루트 노드가 된다.

 

2. 주요 함수들

• heapify(리스트) : 리스트를 힙으로 변환
• heappush(힙, 원소) : 새로운 원소 추가
• heappop(힙) : 힙에서 루트 노드 제거 및 반환
• heapreplace(힙, 원소) : 힙에서 루트 노드를 제거하고 새로운 원소 추가 ( pop&push 한 번에! )
• nlargeset(n, iterable, key=None) : iterable 객체에서 n개의 가장 큰 원소 반환(리스트 형태)
• nsamllest(n, iterable, key=None) : iterable 객체에서 n개의 가장 작은 원소 반환(리스트 형태) 

 

3. 예시

0) 모듈 import 및 리스트 선언

import heapq

li = [9, 4, 2, 8, 5, 1, 7 ,6, 3]

 

1) heapify(리스트)

heapq.heapify(li)

print(li)

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

리스트를 최소힙으로 변환하였다.

가장 작은 원소가 루트 노드 자리인 가장 앞에 위치하고, 부모노드가 항상 자식 노드보다 작은 값을 가진다.

부모노드 인덱스 = (자식노드 인덱스- 1) // 2 (2로 나눈 몫)

 

 

2) heappush(힙, 원소)

heapq.heappush(li, 10)
print(li)

>>> [1, 3, 2, 4, 5, 9, 7, 6, 8, 10]

원소를 push한다. 이때 heap의 성질을 유지하며 push 된다.

 

 

3) heappop(힙)

sItem = heapq.heappop(li)
print(sItem)
print(li)

>>> 1
>>> [2, 3, 7, 4, 5, 9, 10, 6, 8]

힙의 루트노드를 삭제하고 반환한다.

 

 

4) heapreplace(힙, 원소) pop&push 한번에

heapq.heapreplace(li, 11)
print(li)

>>> [3, 4, 7, 6, 5, 9, 10, 11, 8] # 2가 pop 되고 11 push

루트 노드를 pop하고 인자로 전달한 원소를 push한다.

 

 

5) nlargeset(n, iterable, key=None)

largestItems = heapq.nlargest(3, li)

>>> [11, 10, 9]

인자로 전달한 iterable의 가장 큰 n개를  배열의 형태로 반환한다.

 

print(heapq.nlargest(3,[0,4,3,2,1]))

>>> [4, 3, 2]

 

 

 

6) nsamllest(n, iterable, key=None)

smallestItems = heapq.nsmallest(3, li)

>>> [3, 4, 5]

인자로 전달한 iterable의 가장 작은 n개를  배열의 형태로 반환한다.

 

 

---

4. 최대힙으로 변환하기

heapq 모듈은 기본적으로 최소힙을 지원한다.

하지만 '원소에 -를 붙여서 힙을 구성'한 후

다시 한 번 모든 원소를 양수로 바꾼다면 최대힙을 구성할 수 있다.

import heapq

li = [4, 1, 7, 3, 8, 5]
heap = [-num for num in lst]  # 모든 원소를 음수로 변환
heapq.heapify(heap)  # 힙으로 변환

print(-heapq.heappop(heap))  # 음수를 다시 양수로 변환하여 최대값 출력