[Python] heapq 라이브러리 살펴보기

사용 방법

heap = []            # 비어있는 heap 생성
heappush(heap, item) # heap 값을 추가
item = heappop(heap) # heap에서 가장 작은 값을 삭제
item = heap[0]       # 삭제 없이 heap에서 가장 작은 값 보기
heapify(x)           # 리스트를 heap으로 바꾸기. O(n) 소요
item = heapreplace(heap, item) # 가장 작은 값을 삭제 후 새로운 값을 추가 -> 크기 변화 X

 

HeapQ의 특징

• 모든 K에 대하여  a[k] <= a[2*k+1] and a[k] <= a[2*k+2]  만족
• 0번 index부터 시작
• 정렬 시 O(nlogn) 소요

 

함수 살펴보기

 

Heap의 정렬의 유지를 위한 함수

_siftdown

 parent 변수가 root부터 새로운 값이 들어설 자리까지 아래로 이동

def _siftdown(heap, startpos, pos):
    newitem = heap[pos]
    while pos > startpos:
        parentpos = (pos - 1) >> 1
        parent = heap[parentpos]
        if newitem < parent:
            heap[pos] = parent
            pos = parentpos
            continue
        break
    heap[pos] = newitem

 

_siftup

더 작은 child가 위로 올라올 수 있도록 지정

def _siftup(heap, pos):
    endpos = len(heap)
    startpos = pos
    newitem = heap[pos]
    childpos = 2*pos + 1    # 왼쪽 자식 위치
    while childpos < endpos:
        rightpos = childpos + 1
        if rightpos < endpos and not heap[childpos] < heap[rightpos]:
            childpos = rightpos # 더 작은 자식의 인덱스로 childpos에 저장
        
        heap[pos] = heap[childpos] # 더 작은 child가 위로 올라갈 수 있도록 저장
        pos = childpos
        childpos = 2*pos + 1
    # The leaf at pos is empty now.  Put newitem there, and bubble it up
    # to its final resting place (by sifting its parents down).
    heap[pos] = newitem
    _siftdown(heap, startpos, pos)

 

 

Push, Pop, Replace

def heappush(heap, item):
    heap.append(item)
    _siftdown(heap, 0, len(heap)-1)

def heappop(heap): # 가장 작은 값을 삭제
    lastelt = heap.pop()    # heap이 비어있다면 IndexError 발생
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        _siftup(heap, 0)
        return returnitem
    return lastelt

def heappushpop(heap, item): 
	# heappop 방식을 이용한 더 빠른 heap push 방법
    if heap and heap[0] < item:
        item, heap[0] = heap[0], item
        _siftup(heap, 0)
    return item

 

def heapreplace(heap, item):
	# 가장 작은 값 삭제 및 반환 후, 새로운 값 추가
    # heappop() 한 뒤에 heappush()하는 것보다 훨씬 효율적임
    # 특히 heap 크기가 변화하지 않을 때 더욱 좋음
    returnitem = heap[0]    # heap이 비어있다면 IndexError 발생
    heap[0] = item
    _siftup(heap, 0)
    return returnitem

 

 

기타

리스트를 heap을 변형하기

def heapify(x):
    for i in reversed(range(n//2)):
        _siftup(x, i)

 

리스트를 최대 힙(max heap)으로 변형하기

def _heapify_max(x):
    n = len(x)
    for i in reversed(range(n//2)):
        _siftup_max(x, i)

 

최대값 삭제 및 반환

def _heappop_max(heap):
    """Maxheap version of a heappop."""
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        _siftup_max(heap, 0)
        return returnitem
    return lastelt

def _heapreplace_max(heap, item):
    """Maxheap version of a heappop followed by a heappush."""
    returnitem = heap[0]    # raises appropriate IndexError if heap is empty
    heap[0] = item
    _siftup_max(heap, 0)
    return returnitem

 

 

https://docs.python.org/ko/3.10/library/heapq.html

heapq — 힙 큐 알고리즘 — Python 3.10.0rc1 문서

 

 

+) 추가

queue 라이브러리의 PriorityQueue보다 heapq 라이브러리가 더 빠르다.