[Python] heapq 모듈 사용법

 

Python에서는 heap 자료구조를 쉽게 구현할 수 있도록 도와주는 내장 heapq 모듈이 존재한다. heap을 직접 구현하는 것보다 훨씬 편리하고 내장되어 있는 모듈이기에 코딩 테스트를 위해서 사용법을 알아두는 것이 도움이 될 것이다.

 

heapq

import heapq

여기서 중요한 점은 Python에서 heap은 list 기반으로 동작하고, heap의 root가 가장 작은 값을 가지는 최소 힙(min-heap)이다.

 

heapq.heappush

heap에 값을 넣으려면 heappush 메서드를 사용한다. 첫 번째 인수는 heap으로 사용될 list가 들어가고 두 번째 인수로는 넣고자 하는 값이 들어간다.

import heapq

heap = []

# heap.heappush(list, item)
heapq.heappush(heap, 3)
heapq.heappush(heap, 4)
heapq.heappush(heap, 1)

print(heap) # Result [1, 4, 3]

heapq는 min-heap으로 동작하기에 heap[0]의 값이 1 인 것을 알 수 있다.

 

heapq.heappop

heap에서 root 값을 빼내고자 할 때는 heappop 메서드를 사용하고, 첫 번째 인수로 heap list가 들어간다. 그리고 heappop 메서드는 heap에서 가장 작은 값을 리턴한다.

import heapq

heap = []

# heap.heappush(list, item)
heapq.heappush(heap, 3)
heapq.heappush(heap, 4)
heapq.heappush(heap, 1)

print(heap) # Result [1, 4, 3]

# heap.heappop(list)
print(heapq.heappop(heap)) # Result 1

print(heap) # Result [3, 4]

 

heapq.heapify

이미 주어진 list를 heap으로 만들 때는 heapify 메서드를 활용한다.

import heapq

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

# heapq.heapify(list)
heapq.heapify(heap)

print(heap) # Result [1, 2, 3, 6, 5, 4, 7, 8, 9]

 

heapq.nlargest / heapq.nsmallest

heap에서 가장 큰 n개의 elements 혹은 가장 작은 n개의 elements를 찾기 위해 사용하는 메서드이다.

import heapq

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

# heapq.heapify(list)
heapq.heapify(heap)

print(heapq.nsmallest(3, heap)) # Result [1, 2, 3]
print(heapq.nlargest(3, heap)) # Result [9, 8, 7]

 

heapq로 max-heap 구현

기본적으로 heapq 모듈은 min-heap만을 제공하지만, item을 tuple 형태로 넣어주게 된다면 max heap을 구현 할 수 있다.

import heapq

heap = []

for i in range(1, 10):
    heapq.heappush(heap, (-i, i))
    
print(heap) # [(-9, 9), (-8, 8), (-6, 6), (-7, 7), (-3, 3), (-2, 2), (-5, 5), (-1, 1), (-4, 4)]

for _ in range(9):
    print(heapq.heappop(heap)[1]) # Result 9 8 7 6 5 4 3 2 1

(-i)를 기준으로 한 min-heap이 만들어지기 때문에 i를 기준으로는 max-heap이 만들어진다.