python에서 힙은 어떻게 사용할까?

heapq

파이썬에서 힙을 이용하는 이유는 아마 주로 우선순위 큐를 사용하기 위해서일거라고 생각한다. 우선순위 큐는 힙으로 구현하는 것이 성능이 가장 좋고, 이를 파이썬에서도 heapq 모듈을 제공해서 편리하게 우선순위 큐를 사용할 수 있게 한다.

heapq 모듈은 이진 트리(binary tree) 기반의 최소 힙(min heap) 자료 구조를 제공한다. 최소 힙을 사용하면 원소들이 항상 정렬된 상태로 추가, 삭제되고, 최소 힙에서 가장 작은 값은 인덱스 0인 루트에 위치한다. 원소가 추가되고 삭제될 때도 최소 힙의 모든 원소(k)는 자식 원소들(2k+1, 2k+2)보다 크기가 작거나 같도록 원소가 추가되고 삭제된다.

     1  ---> root
   /   \
  3     5
 / \   /
4   10 7

• 원소 추가 : heappush(list, x)
• 원소 삭제 : heappop(list)

import heapq

heap = [] // heapq 모듈은 기본 리스트를 최소 힙처럼 다룰 수 있게 한다. heapq 모듈의 함수를 호출할 때마다 이 리스트를 인자로 넘긴다.

heapq.heappush(heap, 4)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3) // [1, 3, 7, 4]

print(heapq.heappop(heap)) // 1

• 성능
• 원소 추가 : O(logN)
• 원소 삭제 : O(logN)

최솟값 얻기

만든 최소 힙에서 최솟값을 삭제하지 않고 얻는 방법은 리스트의 0번째 인덱스에 최솟값이 위치하므로 바로 접근하면 된다.

print(heap[0])

기존 리스트 힙으로 변환

기존 리스트를 힙으로 변환하려면 heapify() 함수를 사용한다.

heap = [4, 1, 7, 3, 8, 5]
heapq.heapify(heap)
print(heap)

성능 : O(N) (인자로 넘기는 리스트의 원소수에 비례한다)

최대 힙

heapq 모듈은 최소 힙만을 지원하기 때문에 최대 힙처럼 동작하게 하려면 튜플(tuple)을 이용하면 된다. 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 원리를 이용하면 최대 힙을 만들 수 있다.

따라서 최대 힙을 만드려면 각 값에 대한 우선 순위를 구한 후, (우선 순위, 값) 구조의 튜플을 힙에 추가하거나 삭제하면된다. 값을 읽을 때는 튜플에서 인덱스 1에 있는 값을 읽으면 된다.

import heapq

nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
  heapq.heappush(heap, (-num, num))
  
print(heapq.heappop(heap)[1])

K번째 최소/최댓값

최소 힙이나 최대 힙을 사용하면 k번째 최솟값이나 최댓값을 효율적으로 구할 수 있다. 단순히 힙에서 k번만큼 루트를 제거해주면 된다.

import heapq

def kth_smallest(nums, k):
  heap = heapq.heapify(nums)
  
  kth_min = None
  
  for _ in range(k):
    kth_min = heapq.heappop(heap)
   
  return kth_min

힙 정렬

import heapq

def heap_sort(nums):
  heap = heapq.heapify(nums)
  
  sorted_nums = []
  
  while heap:
    sorted_nums.append(heapq.heappop(heap))
  return sorted_nums

출처 : https://www.daleseo.com/python-heapq/