[파이썬] heapq 모듈 코드 뜯어보기

문제 풀다가

백준문제 11279 최대힙 문제는 최대힙을 구현하는 문제입니다. 최대힙을 직접 구현해도 되지만 파이썬은 heapq라는 힙 모듈을 제공합니다.

함수를 이용해 힙을 직접 구현한 코드는 다음과 같습니다. 클래스를 이용해 구현할 수도 있습니다.

import sys

read = sys.stdin.readline


def up_heapify(index, queue):
    child_index = index
    while child_index != 0:
        parent_index = (child_index - 1) // 2
        if queue[parent_index] < queue[child_index]:
            queue[parent_index], queue[child_index] = queue[child_index], queue[parent_index]
            child_index = parent_index
        else:
            return


def find_bigger_child_index(index, heap_size):
    parent = index
    left_child = (parent * 2) + 1
    right_child = (parent * 2) + 2

    if left_child < heap_size and priority_queue[parent] < priority_queue[left_child]:
        parent = left_child
    if right_child < heap_size and priority_queue[parent] < priority_queue[right_child]:
        parent = right_child
    return parent


def down_heapify(index, queue):
    parent_index = index
    bigger_child_index = find_bigger_child_index(parent_index, len(queue))
    while parent_index != bigger_child_index:
        queue[parent_index], queue[bigger_child_index] = queue[bigger_child_index], queue[parent_index]
        parent_index = bigger_child_index
        bigger_child_index = find_bigger_child_index(parent_index, len(queue))


def heap_pop(queue):
    if not len(queue):
        return 0
    tmp = priority_queue[0]
    priority_queue[0] = priority_queue[-1]
    priority_queue.pop()
    down_heapify(0, queue)
    return tmp


N = int(read())
priority_queue = []
for _ in range(N):
    order = int(read())
    if order != 0:
        priority_queue.append(order)
        up_heapify(len(priority_queue) - 1, priority_queue)
    else:
        print(heap_pop(priority_queue))

heapq 모듈을 이용한 코드는 다음과 같습니다.

import sys
import heapq

read = sys.stdin.readline

n = int(read())

heap = []
for i in range(n):
    num = int(read())
    if num == 0:
        if not heap:
            print(0)
        else:
            print(heapq.heappop(heap)[1])
    else:
        heapq.heappush(heap, (-num, num))

그런데 몇번 돌려봤지만 작동시간에 거의 차이가 나지 않았습니다. 어떻게 된걸까요?

파이썬 heapq 모듈 뜯어보기

heapq 모듈이 어떻게 되어있는지 파악하기 위해 모듈을 뜯어보았습니다. 아래는 해당 소스입니다.

https://github.com/python/cpython/blob/master/Lib/heapq.py

heapq 모듈 역시 배열을 이용하여 heap을 구현하고 있었습니다. heappush, heappop 함수를 보면 _siftdown, _siftup이라는 함수를 사용합니다. 해당 부분이 down_heap, up_heap 부분이라고 생각됩니다.

heapq 모듈역시 while 문을 사용하여 down heap, up heap을 구현하고 있었습니다. 이렇게 직접 작성한 코드와 모듈이 비슷한 로직을 가지고 있기 때문에 유사한 시간이 나올 수 있었다고 생각됩니다.

배운점

코드를 직접 뜯어봤을 때 작동 원리가 직접 구현한 방법과 비슷했던건 신기한 경험이었습니다. 어떤 라이브러리를 사용할 때, 해당 라이브러리가 어떤 방식으로 구동되는지 확인하며 얻는점이 많은 것 같습니다. 공부하면서 소스코드를 뜯어보는 경험을 통해 나중에 오픈소스 기여에 발판이 될 습관을 만들어 둬야겠습니다.