백준 문제 풀이 - 최소 힙 1927번

📜 문제 이해하기

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

💡 문제 재정의

최소 힙을 구성한 뒤 최소값부터 차례대로 출력해보자.

✏️ 계획 수립

최소 힙은 다음의 조건을 가진다.
1. 부모 노드는 자식 노드보다 항상 작다.
2. 마지막 레벨을 제외한 모든 레벨에 노드가 꽉 차있어야 한다.
3. 마지막 레벨에 노드가 있을때는 항상 가장 왼쪽부터 순서대로 채워져 있어야 한다.

힙을 구성하는 배열이 존재할 때 다음의 값을 만족한다.
A[i]에 왼쪽 자손은 A[2*i+1]
A[i]에 오른쪽 자손은 A[2*i+2]
A[i]에 부모는 A[(i-1)/2]

힙에 값을 넣을 insert 함수는 다음과 같이 구현된다.
1. 힙 배열에 마지막 원소에 값을 추가한다.
2. 부모 원소와 비교한뒤 부모 원소가 자식 보다 크다면 서로 위치를 바꾸고 현재 idx를 부모 idx로 변경한다.
3. 위 과정을 idx 0까지 반복한다.
4. 만약 부모 원소가 현재 원소보다 작다면 종료한다.

힙에 값을 뺄 pop 함수는 다음과 같이 구현된다.
1. 가장 앞에 원소를 빼고 가장 뒤의 원소를 앞에 원소로 교체한다.
2. 위의 원소에서 왼쪽 노드와 비교해보고 왼쪽 노드가 더 작다면 왼쪽노드와 교체한다.
3. 만약 오른쪽 노드가 존재하고 왼쪽 노드보다 오른쪽 노드가 더 작다면 오른쪽 노드와 교체한다.
4. 위 과정을 적절한 위치에 내려올때까지 반복한다.

💻 계획 수행

import sys


class Heap:
    def __init__(self):
        self.min_heap = []

    def insert(self, key):
        self.min_heap.append(key)
        idx = len(self.min_heap) - 1
        while idx > 0 and self.min_heap[int((idx - 1) / 2)] > self.min_heap[idx]:
            tmp = self.min_heap[int((idx - 1) / 2)]
            self.min_heap[int((idx - 1) / 2)] = self.min_heap[idx]
            self.min_heap[idx] = tmp
            idx = int((idx - 1) / 2)

    def pop_heap(self):
        if len(self.min_heap) > 1:
            min_value = self.min_heap[0]
            self.min_heap[0] = self.min_heap.pop()
            here = 0
            length = len(self.min_heap)
            while True:
                left = here * 2 + 1
                right = here * 2 + 2
                if left >= length:
                    break

                next_idx = here
                # 부모와 왼쪽 자식과 비교하고 왼쪽 자식이 더 작다면 넣기
                if self.min_heap[here] > self.min_heap[left]:
                    next_idx = left
                # 만약 오른쪽 자식이 있다면 왼쪽 자식과 비교후에 오른쪽 자식이 더 작다면 바꾸기
                if right < length and self.min_heap[next_idx] > self.min_heap[right]:
                    next_idx = right

                if next_idx == here:
                    break

                tmp = self.min_heap[next_idx]
                self.min_heap[next_idx] = self.min_heap[here]
                self.min_heap[here] = tmp

                here = next_idx
            return min_value
        else:
            return self.min_heap.pop()


if __name__ == '__main__':
    heap = Heap()
    N = int(sys.stdin.readline())
    result = []
    for _ in range(N):
        cmd = int(sys.stdin.readline())
        if cmd != 0:
            heap.insert(cmd)
        else:
            if len(heap.min_heap) > 0:
                result.append(heap.pop_heap())
            else:
                result.append(0)

    print(*result, sep='\n')

🤔 회고

파이썬의 모듈 heapq를 사용하면 간단하게 해결할 수 있는 문제지만 공부를 위해서 직접 구현해 보았다.