[Python] 힙(Heap) 구현하기

힙?

힙은 최댓값/최솟값을 빠르게 찾을 수 있는 완전 이진 트리 기반의 자료구조이다.

• 최소 힙(Min Heap) - 최솟값이 항상 루트에 위치함
• 최대 힙(Max Heap) - 최댓값이 항상 루트에 위치함

Min Heap 구현하기

파이썬에서는 heapq 모듈을 사용해 힙을 쉽게 구현할 수 있다.

import heapq

heap = []

heapq.heappush(heap, 5)
heapq.heappush(heap, 2)
heapq.heappush(heap, 8)
heapq.heappush(heap, 1)

print(heapq.heappop(heap))

리스트를 이용해 힙을 생성한 이후,
heappush(힙, 요소) 를 통해 요소를 추가하고
heappop(힙) 을 통해 루트 요소를 반환하면서 제거 가능하다.

Max Heap 구현하기

heapq 라이브러리는 기본적으로 Min Heap으로 동작한다. 만약 Max Heap을 구현하고 싶다면 heappush로 요소 저장 시 -를 붙여주면 된다.

힙에서 꺼낼 때 다시 -를 붙여 원래 값으로 변환한다.

heapq.heappush(heap, -5)
heapq.heappush(heap, -2)
heapq.heappush(heap, -8)
heapq.heappush(heap, -1)

print([-x for x in heap])

print(-heapq.heappop(heap))  
print([-x for x in heap]) 

리스트를 힙으로 변환하기

heapify() 함수를 이용하면 리스트를 한번에 힙으로 변환시켜준다.

가장 작은 n개 원소 찾기

heapq.nsmallest(n, 리스트) - 가장 작은 n개 값 반환
heapq.nlargest(n, 리스트) - 가장 큰 n개 값 반환