본문 바로가기

Python

[Python] heapq (min heap)

힙은 모든 부모 노드가 자식보다 작거나 같은 값을 갖는 이진 트리이다. 힙은 모든 k에 대해 heap[k] <= heap[2*k+1] heap[k] <= heap[2*k+2]인 배열을 사용한다. 요소는 0부터 센다. 비교를 위해, 존재하지 않는 요소는 무한으로 간주한다. 힙의 가장 작은 요소는 항상 루트인 heap[0]이다.

 

힙을 만들려면, []로 초기화된 리스트를 사용하거나, 함수 heapify()를 통해 값이 들어 있는 리스트를 힙으로 변환 할 수 있다.  

함수

heapq.heappush(heap, item)

힙 불변성을 유지하면서, item 값을 heap으로 푸시한다.

 

heapq.heappop(heap)

힙 불변성을 유지하면서, heap에서 가장 작은 항목을 팝하고 반환한다. 힙이 비어 있으면, IndexError가 발생한다. 팝 하지 않고 가장 작은 항목에 액세스하려면, heap[0]을 사용하면 된다.

 

heapq.heappushpop(heap, item)

힙에 item을 푸시한 다음, heap에서 가장 작은 항목을 팝하고 반환한다. 결합한 액션은 heappush()한 다음 heappop()을 별도로 호출하는 것보다 더 효율적으로 실행한다.

 

heapq.heapify(x)

리스트 x를 선형 시간으로 제자리에서 힙으로 변환한다.

 

heapq.heapreplace(heap, item)

heap에서 가장 작은 항목을 팝하고 반환하며, 새로운 item도 푸시한다. 힙 크기는 변경되지 않는다. 힙이 비어 있으면, IndexError가 발생한다. 

이 한 단계 연산은 heappop()한 다음 heappush()하는 것보다 더 효율적이며 고정 크기 힙을 사용할 때 더 적합 할 수 있다. 팝/푸시 조합은 항상 힙에서 요소를 반환하고 그것을 item으로 대체한다.

반환된 값은 추가된 item보다 클 수 있다. 그것이 바람직하지 않다면, 대신 heappushpop() 사용하면 된다. 푸시/팝 조합은 두 값 중 작은 값을 반환하여, 힙에 큰 값을 남겨 둔다.

 

heapq.merge(*iterables, key=None, reverse=False)

여러 정렬된 입력을 단일 정렬된 출력으로 병합한다 (예를 들어, 여러 로그 파일에서 타임 스탬프 된 항목을 병합한다). 정렬된 값에 대한 이터레이터를 반환한다.

sorted(itertools.chain(*iterables))와 비슷하지만 이터러블을 반환하고, 데이터를 한 번에 메모리로 가져오지 않으며, 각 입력 스트림이 이미 (최소에서 최대로) 정렬된 것으로 가정한다.

키워드 인자로 지정해야 하는 두 개의 선택적 인자가 있다.

key는 각 입력 요소에서 비교 키를 추출하는 데 사용되는 단일 인자의 키 함수를 지정한다. 기본값은 None이다 (요소를 직접 비교한다).

reverse는 불리언 값이다. True로 설정하면, 각 비교가 반대로 된 것처럼 입력 요소가 병합된다. sorted(itertools.chain(*iterables), reverse=True)와 유사한 동작을 달성하려면 모든 이터러블이 최대에서 최소로 정렬되어 있어야 한다.

 

heapq.nlargest(n, iterable, key=None)

iterable에 의해 정의된 데이터 집합에서 n개의 가장 큰 요소로 구성된 리스트를 반환한다. key가 제공되면 iterable의 각 요소에서 비교 키를 추출하는 데 사용되는 단일 인자 함수를 지정한다 (예를 들어, key=str.lower). 다음과 동등하다다: sorted(iterable, key=key, reverse=True)[:n].

 

heapq.nsmallest(n, iterable, key=None)

iterable에 의해 정의된 데이터 집합에서 n개의 가장 작은 요소로 구성된 리스트를 반환한다. key가 제공되면 iterable의 각 요소에서 비교 키를 추출하는 데 사용되는 단일 인자 함수를 지정한다 (예를 들어, key=str.lower). 다음과 동등하다: sorted(iterable, key=key)[:n].

 

마지막 두 함수는 작은 n 값에서 가장 잘 동작한다. 값이 크면, sorted() 기능을 사용하는 것이 더 효율적이다. 또한, n==1일 때는, 내장 min() max() 함수를 사용하는 것이 더 효율적이다. 이 함수를 반복해서 사용해야 하면, iterable을 실제 힙으로 바꾸는 것이 좋다.

예시

from heapq import *

h = []
heappush(h, (2, 'orange'))
heappush(h, (0, 'apple'))
heappush(h, (4, 'banana'))
heappop(h) # (0, 'apple')

[응용] 최대 힙

heapq 모듈은 최소 힙(min heap) 기능만을 제공하기 때문에 최대 힙(max heap)으로 활용하려면 약간의 요령이 필요하다. 바로 힙에 튜플(tuple)로 원소를 추가허가나 삭제하면, 튜퓰 내에서 맨 앞에 있는 값을 기준으로 최소힙이 구성되는 원리를 이용하는 것이다.

따라서 최대 힙을 만들려면 각 값에 대한 우선 순위를 구한 후, (우선순위, 값) 구조의 튜플(tuple)을 힙에 추가하거나 삭제하면 된다. 그리고 힙에서 값을 읽어올 때는 각 튜플에서 인텍스 1에 있는 값을 취하면 된다.

예시

import heapq

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

for num in nums:
  heapq.heappush(heap, (-num, num))  # (우선 순위, 값)

while heap:
  print(heapq.heappop(heap)[1])  # index 1

결과

8
7
5
4
3
1
참고: https://docs.python.org/ko/3/library/heapq.html,
https://www.daleseo.com/python-heapq/