[Python] 힙(heap) 제공 라이브러리 heapq

heapq

힙(Heap) 기능을 제공하기 위한 라이브러리이다.

파이썬의 힙은 최소 힙(Min Heap)으로 구성되어 있으므로 단순히 원소를 힙에 전부 넣었다가 뺴는 것만으로도 시간 복잡도 O(NlogN)에 오름차순 정렬이 완료된다. (보통 최소 힙 자료구조의 최상단 원소 : 가장 작은 원소)

 

heapq.heappush( ) : 힙에서 원소를 삽입할 때

heapq.heappop( ) : 힙에서 원소를 꺼낼 때 항상 가장 작은 원소를 꺼내게 된다. 

 

예) 힙 정렬(Heap Sort)을 heapq로 구현하는 경우

import heapq

def heapsort(iterable):
    h = []
    result = []
    
    # 모든 원소를 차례대로 힙에 삽입
    for value in iterable:
        heapq.heappush(h, value)
        
    # 힙에 삽입된 모든 원소를 차례대로 꺼내어 담기
    for i in range(len(h)):
        result.append(heapq.heappop(h))
        
    return result
    
result = heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
print(result) # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

 

파이썬에서는 최대 힙(Max Heap)을 제공하지 않는다. 따라서 heapq 라이브러리를 이용해 최대 힙을 구현해야하는 데, 이때는 원소의 부호를 임시로 변경하는 방식을 사용한다. → 힙에 원소를 삽입하기 전에 잠시 부호를 반대로 바꾸었다가, 힙에서 원소를 꺼낸 뒤에 다시 원소의 부호를 바꾼다.

 

예) 최대 힙을 heapq로 구현하여 내림차순 힙 정렬을 구현하는 경우

import heapq

def heapsortDesc(iterable):
    h = []
    result = []
    
    # 모든 원소를 차례대로 힙에 삽입(부호를 변경하여)
    for value in iterable:
        heapq.heappush(h, -value)
        
    # 힙에 삽입된 모든 원소를 차례대로 꺼내어 담기
    for i in range(len(h)):
        result.append(-heapq.heappop(h))
        
    return result
    
result = heapsortDesc([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
print(result) # [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

 

코딩 테스트>> heapq는 다익스트라 최단 경로 알고리즘을 포함해 다양한 알고리즘에서 우선순위 큐 기능을 구현하고자 할 때 사용된다. heapq 외에도 priorityQueue 라이브러리를 사용할 수 있지만, 코딩 테스트 환경에서는 보통 heapq가 더 빠르게 동작한다. 

 

Ref.

이것이 취업을 위한 코딩 테스트다 with 파이썬