[Python] 백준 1927 - 최소 힙(Min Heap)

문제

https://www.acmicpc.net/problem/1927

 

해설

최소 힙(Min Heap)을 이용하여 해결하는 문제였다. Heap은 기본적으로 이진트리 구조로 형성되므로 insert와 delete 연산의 시간 복잡도는 \( O(logn) \)이 된다. 여기서 \( log \)는 일반적으로 밑이 2인 로그로 \( log_{2}{n} \)를 얘기한다. 이 문제에서는 입력 횟수만큼 반복하므로 \( O(Nlogn) \)이 된다.

 

Fig 1. Min Heap

 

패키지를 사용하여 해결하는 방법과 직접 구현 두 가지 방식으로 풀어보았다. 이번에도 input() 메소드가 시간초과를 유발하여 PyPy3를 이용하였다.

 

Solution 1.

아래는 heap을 이용하여 해결한 방식이다.

 

# PyPy3
import heapq


def solution(lst):

    heap = []
    for c in lst:
        if c == 0:
            if heap:
                print(heapq.heappop(heap))
            else:
                print(0)
        else:
            heapq.heappush(heap, c)


n = int(input())
node = [int(input()) for _ in range(n)]
solution(node)

 

Solution 2.

아래는 패키지를 사용하지 않고 직접 heap을 구현해서 풀어보았다.

 

# PyPy3
class MinHeap:
    def __init__(self):
        self.heap = []

    def insert(self, val):
        # 값을 힙에 추가하고, 최소 힙의 특성을 유지하도록 함
        self.heap.append(val)
        self._heapify_up(len(self.heap) - 1)

    def extract_min(self):
        if not self.heap:
            return None
        # 최솟값은 항상 루트 노드에 있음
        min_val = self.heap[0]
        # 마지막 요소를 루트로 옮긴 후 최소 힙을 유지
        self.heap[0] = self.heap[-1]
        self.heap.pop()
        self._heapify_down(0)
        return min_val

    def peek(self):
        return self.heap[0] if self.heap else None

    def size(self):
        return len(self.heap)

    def _heapify_up(self, index):
        # 부모 노드와 비교하여 최소 힙을 유지하도록 위치 변경
        parent_index = (index - 1) // 2
        if index > 0 and self.heap[index] < self.heap[parent_index]:
            self.heap[index], self.heap[parent_index] = self.heap[parent_index], self.heap[index]
            self._heapify_up(parent_index)

    def _heapify_down(self, index):
        # 자식 노드들과 비교하여 최소 힙을 유지하도록 위치 변경
        left = 2 * index + 1
        right = 2 * index + 2
        smallest = index

        if left < len(self.heap) and self.heap[left] < self.heap[smallest]:
            smallest = left
        if right < len(self.heap) and self.heap[right] < self.heap[smallest]:
            smallest = right
        if smallest != index:
            self.heap[index], self.heap[smallest] = self.heap[smallest], self.heap[index]
            self._heapify_down(smallest)


def solution(lst):

    heap = MinHeap()
    for c in lst:
        if c == 0:
            if heap.size():
                print(heap.extract_min())
            else:
                print(0)
        else:
            heap.insert(c)


n = int(input())
node = [int(input()) for _ in range(n)]
solution(node)

 

---

 

참고 자료

 

https://discourse.opengenus.org/t/max-heap-and-min-heap/3388