[Python] 힙 자료구조 / heapq

최근에 이것이 취업을 위한 코딩테스트다 with Python에서 2019년도 카카오 그리디 알고리즘 문제를 푸는데 우선순위 큐(heapq) 라이브러리를 몰라서 많이 어려웠던 경험이 있었다.

이번 포스팅에서는 heap 자료구조는 무엇이고, 파이썬의 heapq 모듈을 사용하는 방법에 대해서 알아보려고 한다.

들어가기 전에..

우선순위 큐: 우선순위의 개념을 큐에 도입한 자료구조이다.
데이터들이 우선순위를 가지고 있고 우선순위가 높은 데이터가 먼저 나간다.

우선순위 큐는 배열, 연결리스트, 힙으로 구현이 가능하다. 이 중에서 힙으로 구현하는 것이 가장 효율적이다. (삽입과 삭제가 모두 동일하게 O(logN))

힙이란?

힙은 특정한 규칙을 갖는 트리로, 최댓값과 최솟값을 찾는 연산을 빠르게 하기 위해 고안된 완전이진트리를 기본으로 한다.

힙 property: A가 B의 부모노드이면 A의 키값과 B의 키값 사이에는 대소 관계가 성립한다.

• 최소 힙(Min Heap): 부모 노드의 키값이 자식 노드의 키값보다 항상 작은 heap.
• 최대 힙(Max Heap): 부모 노드의 키값이 자식 노드의 키값보다 항상 큰 heap.
  <출처: https://www.geeksforgeeks.org/heap-data-structure/minheapandmaxheap/>

이러한 속성으로 인해서 힙에서는 가장 낮은(혹은 높은) 우선순위를 가지는 노드가 항상 루트에 오게 되고 이를 이용해서 우선순위 큐와 같은 추상적 자료형을 구현할 수 있다.
(이 때, 키값의 대소 관계는 부모/자식 간에만 성립하고 형제노드 사이에는 성립하지 않는다)

Python heap 자료구조

• 파이썬 heapq 모듈은 heapq(Priority Queue) 알고리즘을 제공한다.
• 모든 부모 노드는 그의 자식 노드보다 값이 작거나 큰 이진트리 구조인데, 내부적으로는 인덱스 0에서 시작해서 k번째 원소가 항상 자식 원소들 (2k+1, 2k+2)보다 작거나 같은 최소 힙의 형태로 정렬된다.

(1) 함수 사용하기

• heapq.heappush(heap, item): item을 heap에 추가
• heapq.heappop(heap): heap에서 가장 작은 원소를 pop후 return. 비어있는 경우에는 IndexError가 호출된다.
• heapq.heapify(x): 리스트 x를 즉각적으로 heap으로 변환시킨다.

(2) heap 생성 및 원소 추가

heapq는 리스트를 Min-Heap처럼 다룰 수 있도록 하기 때문에 빈 리스트를 생성한 후에 heapq의 함수를 호출할 때마다 리스트를 인자에 넘겨야 한다.

import heapq

heap = []
heapq.heappush(heap, 40)
heapq.heappush(heap, 10)
heapq.heappush(heap, 20)

print(heap)   # [10, 40, 20]

만약 이미 생성해둔 리스트가 있으면 heapify를 통해 바로 heap으로 변경 가능하다.

heap2 = [40, 10, 20]
heapq.heapify(heap2)

print(heap2)   # [10, 40, 20]

(3) heap 원소 삭제

heappop 함수는 가장 작은 원소를 힙에서 제거함과 동시에 결과값으로 return.

result = heapq.heappop(heap)

print(result)   # 10
print(heapq)   # [20, 40]

heap에서 가장 작은 원소인 10이 return 되었고, 힙에서 제거되었다.
제거를 한 이후에도 Min-Heap 구조를 만족시키고 있었다.

(4) 심화) Max-Heap 만들기

기본적으로 파이썬의 heapq 모듈은 Min-Heap으로 구현이 되어 있기 때문에 Max-Heap 구현을 위해서는 약간의 트릭이 필요하다.
=> y=-x 변환을 하게 되면 최솟값 정렬이 최댓값 정렬로 바뀐다!!

힙에 원소를 추가할 때 (-item, item) 튜플 형태로 넣어주면 튜플의 첫 번째 원소로 우선순위 힙을 구성하게 되는데 원소값의 부호가 바뀌었기 때문에 최소 힙으로 구현된 heapq 모듈을 최대 힙 구현에 활용하게 되는 것이다.

heap_items = [1,3,5,7,9]

max_heap = []
for item in heap_items:
  heapq.heappush(max_heap, (-item, item))

print(max_heap)  # [(-9, 9), (-7, 7), (-3, 3), (-1, 1), (-5, 5)]

max_item = heapq.heappop(max_heap)[1]
print(max_item)   # 9

heappop을 사용하게 되면 힙에 있는 최댓값이 반환되는 것을 확인할 수 있다. 이 때 실제 원소 값은 튜플의 두 번째 자리에 저장되어 있다.

참고

https://docs.python.org/2/library/heapq.html
https://littlefoxdiary.tistory.com/3
https://gmlwjd9405.github.io/2018/05/10/data-structure-heap.html