[알고리즘] heapq모듈을 이용한 최소 힙 문제풀이

이틀 전에 최대 힙을 공부했고, 스택이나 큐와 같은 알고리즘에 익숙해졌기 때문에

오늘은 백준에 있는 힙 관련 문제중에 최소 힙 문제를 풀어보려고 한다.

백준 1927번 문제 - 최소  힙 

문제 설명 널리 잘 알려진 자료구조 중 최소 힙이 있다. 최소 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오. 배열에 자연수 x를 넣는다. 배열에서 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다. 프로그램은 처음에 비어있는 배열에서 시작하게 된다.

문제를 보고 처음에 든 생각은 최대 힙 알고리즘을 가져와서 반대로 진행시키면 된다는 생각으로

배웠던 알고리즘 코드를 참고하여  코드를 구현했다.

 

※ 처음 문제풀이

import sys

n = int(input()) # ex) 9
arr = [None]
for _ in range(n):
    x = int(sys.stdin.readline()) # ex) 0, 12345678, 1, 2, 0, 0, 0, 0, 32
    if x == 0: # x가 0이면 가장 작은 수를 pop()한다.
        if len(arr)-1 < 1: # 하지만 arr에 None밖에 없으면 0을 출력한다.
            print(0)
        else:
            arr[1], arr[-1] = arr[-1], arr[1]
            print(arr.pop())
            #이 상태에서 arr[1]이 다시 배열의 최솟값이 되도록 해야함.

    else:
        arr.append(x)
        # 최소 힙 구현, x를 추가했을때 root가 되게 해야함
        cur = len(arr) - 1
        parent = cur // 2
        while parent > 0:
            if arr[cur] < arr[parent]:
                arr[parent], arr[cur] = arr[cur], arr[parent]
            cur = parent
            parent = cur // 2

우선 input부터 시작해서 이번에도 최대 입력 횟수가 100000이 될 수 있기 때문에 x를 sys.stdin.readline()으로 받았고,

명령이 0일 때 가장 작은 수를 pop(), 그 외에는 배열에 넣고 가장 작은 수를 루트 노드로 만들었다.

하지만 pop()부분을 구현하는데 있어서 pop()을 한 이후에 다시 최솟값을 루트 노드로 만들어야 했는데

이 부분이 생각나지 않아서 강의에 있던 함수를 가져와서 사용했다.

 

※ 최종코드

import sys


def percolate_down(arr, cur):
    smallest = cur
    left = 2 * cur
    right = 2 * cur + 1

    if left <= len(arr) - 1 and arr[left] < arr[smallest]:
        smallest = left

    if right <= len(arr) - 1 and arr[right] < arr[smallest]:
        smallest = right

    if smallest != cur:
        arr[cur], arr[smallest] = arr[smallest], arr[cur]
        percolate_down(arr, smallest)


n = int(input()) # ex) 9
arr = [None]
for _ in range(n):
    x = int(sys.stdin.readline()) # ex) 0, 12345678, 1, 2, 0, 0, 0, 0, 32
    if x == 0: # x가 0이면 가장 작은 수를 pop()한다.
        if len(arr)-1 < 1: # 하지만 arr에 None밖에 없으면 0을 출력한다.
            print(0)
        else:
            arr[1], arr[-1] = arr[-1], arr[1]
            print(arr.pop())
            percolate_down(arr,1)

    else:
        arr.append(x)
        # 최소 힙 구현, x를 추가했을때 root가 되게 해야함
        cur = len(arr) - 1
        parent = cur // 2
        while parent > 0:
            if arr[cur] < arr[parent]:
                arr[parent], arr[cur] = arr[cur], arr[parent]
            cur = parent
            parent = cur // 2

한 번 배운 알고리즘이지만 바로 적용해서 문제풀이를 하는데에는 어려움이 있었다.

하지만 python에는 heapq모듈이 내장되어 있는데 이 모듈은 최소 힙 알고리즘의 구현을 제공해준다.

바로 heapq모듈을 공부하고 이 문제를 풀어보았다.

 

※ heapq모듈 사용

import sys
import heapq

n = int(input())
heap = []
for _ in range(n):
    x = int(sys.stdin.readline())
    if x == 0:
        if not heap:
            print(0)
        else:
            print(heapq.heappop(heap))
    else:
        heapq.heappush(heap, x)

여기서 사용한 heapq모듈의 함수는 두 가지가 있다.

heapq.heappush(heap, item): item값을 heap으로 push하고 최소 힙 구조를 만족한다.
heapq.heappop(heap): 가장 작은 항목을 pop하고 반환한다. 힙이 비어있다면 IndexError가 발생한다.

 

놀라울 정도로 코드가 줄어들었을 뿐만 아니라 heap의 구조를 이해한다면 보기에도 좋다.

다만 문제를 풀면서 느낀점은 알고리즘의 구현을 모듈없이 할 수 있어야 잘 이해했다고 생각하기 때문에,

위의 코드를 리팩토링하여 이해하기 좋고 깔끔한 코드로 만들 것이다.

---

추가적으로 heapq모듈이 포함하고 있는 함수는 이런 것들도 있다.

 

heapq.heapify(x): 리스트 x를 힙으로 변환한다.
heapq.heappushpop(heap, item): 힙에 item을 push한 다음 heap에서 가장 작은 항목을 pop하고 반환한다.
heappush()후 heappop()을 호출하는 것보다 효율적으로 실행된다.
heapq.heapreplace(heap, item): heappushpop이랑은 반대로 pop후에 push한다. 힙이 비어있다면 IndexError가 발생한다.
heapq.nlargest(n, iterable, key=None): 정의된 데이터의 집합에서 n개의 가장 큰 요소로 구성된 리스트를 반환한다.
heapq.nsmallest(n, iterable, key=None): 정의된 데이터의 집합에서 n개의 가장 작은 요소로 구성된 리스트를 반환한다.