[Python] 1927 - 최소 힙 (실버2) : 최소 힙

1. 문제 설명

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

1927번: 최소 힙

 

---

2. 아이디어

이 문제는 아이디어가 필요하다기 보다는 힙(Heap)에 대해 공부해보려고 직접 구현하여 풀었다.

문제 풀이가 목적이라면 그냥 heapq import시켜서 간단하게 해결할 수 있는 문제다.

난이도가 올라갈 수록 인접 행렬보다는 인접 리스트로 풀어야 하는 문제가 많이 나오기 때문에

heap에 대한 이해는 정말 기초 중에 기초로 익혀두어야 한다.

 

---

최소힙

heap은 원소들이 항상 정렬된 상태로 추가/삭제 되는데,

최소힙은 가장 작은 값이 언제나 root(0번째 인덱스)에 위치한다.

하위에 속한 모든 원소(k)는 언제나 자식 원소 (2k+1, 2k+2)보다 작아야 한다.

heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2]

원래 혼자 공부하면서 쓰던 것들이라..ㅎㅎ 악필 양해 부탁드립니다.

입력값으로 5, 12, 1, 2, 3, 46, 55, 12, 1 을 차례로 입력받을 때, 데이터가 최소힙에 저장되는 로직을

간략하게 시각적으로 그려본 것인데 heap에 push하기 위해 필요한 것은 무엇일까.

실제로 값이 2차원으로 배열하는 것이 아니기 때문에 1차원 리스트(혹은 배열)를 이진 트리로

생각하며 인덱스를 나누어야 한다.

 

처음 5를 입력받으면 리스트의 0번째 인덱스, 루트에 위치하게 된다.

그 다음 12라는 값을 push하면 일단 추가해놓고 자신의 부모 원소를 찾아가서 값을 비교해야 한다.

root의 자신 원소의 인덱스는 언제나 1, 2번 이다. 그리고 1번 원소의 자식 인덱스는 언제나 3, 4이고

2번 원소의 자식 인덱스는 5, 6이 된다.

여기서 한 번 더 들어가보면 3번 원소의 자식 인덱스는 7, 8이 된다.

부모 원소를 찾는다는 것은 이 과정을 역행하면 된다. 아니 그럼 이걸 어떻게 하는데??

 

우선 현재 추가된 데이터가 위치한 인덱스를 어찌 잘 알아냈다고 치자.

7번 인덱스에 위치한 데이터라고 가정하면 부모 원소인 3번 인덱스로 접근을 해야하는데,

처음 힙을 공부하면 굉장히 난처하다. 뭔가 수학적으로 접근이 될 거 같은데 쉽게 안 떠오른다.

하지만 우린 이에 대한 해답을 가장 처음에 정의하고 시작했다.

바로 원소 k의 자식 원소는 2k+1, 2k+2라는 것. (직접 이진트리를 보면서 끄적여보면 이해 가능)

 

그럼 자식 위치와 부모 위치를 알았으니 값을 비교해서 정렬해주면 push 기능은 끝난다.

중요한 건 이걸 한 번만 하고 끝낼 것이 아니라 부모 원소보다 값이 큰 경우까지 반복시켜주어야 한다는 것.

---

pop의 기능도 여기서는 단순하다. 가장 작은 값을 출력해달라고 했으니 root의 원소를 빼내면 끝난다.

다만 그렇다고 해서 곧바로 머리를 쳐내면 남아있는 가지들을 재정렬할 방도가 없으니

왕위를 어떻게 잘 물려주고 가야하지 않겠는가. 그래서 일단 가장 마지막 원소와 root원소를 치환하여

마지막 원소를 pop시키고 그 상태에서 재정렬을 시켜주어야 한다.

 

현재 root 원소가 가장 큰 값이 된 상황이니 재정렬의 경우엔 자식 원소와 값의 비교를 한다.

왼쪽 자식과 비교를 하고, 왼쪽 자식은 오른쪽 자식과 또 한 번 비교하여 더 작은 값이 부모 원소로 올라와주어야

나중에 상속 이슈가 발생하지 않는다. (그냥 웃자고 쓴 표현이다.)

이 과정을 반복하면..pop 기능도 사실 끝이다. ㅎ 뭔가 허무하네.

 

---

3. 코드

import sys
input = sys.stdin.readline

class Heap:
    def __init__(self):
        self.data = []
    
    def heappush(self, value):
        self.data.append(value)
        current_index = len(self.data) - 1
        while current_index >= 0:
            if current_index % 2 == 0:
                parent_index = (current_index // 2) - 1
            else:
                parent_index = (current_index // 2)
            if parent_index >= 0 and self.data[parent_index] > self.data[current_index]:
                self.data[current_index], self.data[parent_index] = self.data[parent_index], self.data[current_index]
                current_index = parent_index
            else:
                break
    
    def heappop(self):
        if len(self.data) < 1:
            return 0
        self.data[0], self.data[-1] = self.data[-1], self.data[0]
        deleted_value = self.data.pop()
        self.re_sort(0)
        return deleted_value

    def re_sort(self, idx):
        current = idx
        left = idx * 2 + 1
        right = idx * 2 + 2

        if left <= len(self.data) - 1 and self.data[idx] > self.data[left]:
            current = left
            if right <= len(self.data) -1 and self.data[left] > self.data[right]:
                current = right
        elif right <= len(self.data) - 1 and self.data[idx] > self.data[right]:
            current = right

        if not current == idx:
            self.data[idx], self.data[current] = self.data[current], self.data[idx]
            self.re_sort(current)

def solution():
    n = int(input())
    heap = Heap()
    for _ in range(n):
        x = int(input())
        if x == 0:
            print(heap.heappop())
        else:
            heap.heappush(x)

solution()

 

파이썬에서 제공하는 heapq를 이용한 풀이

import heapq
import sys
input = sys.stdin.readline

def solution():
    n = int(input())
    heap = []
    for _ in range(n):
        x = int(input())
        if x == 0:
            try:
                print(heapq.heappop(heap))
            except:
                print(0)
        else:
            heapq.heappush(heap, x)
solution()