[알고리즘 / 백준] 1927 - 최소 힙

www.acmicpc.net/problem/1927

1927번: 최소 힙

 

문제는 위와 같으며 최소 힙을 구현할 수 있으면 쉽게 풀 수 있는 문제입니다. 또한 파이썬에서는 heapq 라이브러리를 통해 문제를 보다 쉽게 실수 없이 해결할 수 있습니다.

 

먼저 파이썬에서 heapq 를 사용한 코드를 살펴보면 다음과 같습니다.

from sys import stdin
import heapq

n = int(stdin.readline())
heap = []
for _ in range(n):
    num = int(stdin.readline())

    if num == 0:
        if heap:
            print(heapq.heappop(heap))
        else:
            print(0)
    else:
        heapq.heappush(heap, num)

 

다음으로 파이썬에서 직접 heap 구조를 구현한 코드는 다음과 같습니다.

from sys import stdin

class Heap:
    def __init__(self):
        self.heap = list()
        self.heap.append(None)

    def move_up(self, idx):
        if idx <= 1:
            return False

        parent_idx = idx // 2
        if self.heap[idx] < self.heap[parent_idx]:
            return True
        else:
            return False

    def move_down(self, idx):
        left_child_idx = idx * 2
        right_child_idx = idx * 2 + 1

        # Case1: 왼쪽 자식도 없을 때
        if left_child_idx >= len(self.heap):
            return False

        # Case2: 오른쪽 자식만 없을 때
        elif right_child_idx >= len(self.heap):
            if self.heap[idx] > self.heap[left_child_idx]:
                return True
            else:
                return False

        # Case3: 왼쪽, 오른쪽 자식 모두 있을 때
        else:
            if self.heap[left_child_idx] < self.heap[right_child_idx]:
                if self.heap[idx] > self.heap[left_child_idx]:
                    return True
                else:
                    return False
            else:
                if self.heap[idx] > self.heap[right_child_idx]:
                    return True
                else:
                    return False

    def insert(self, data):
        self.heap.append(data)

        idx = len(self.heap) - 1
        while self.move_up(idx):
            parent_idx = idx // 2
            self.heap[idx], self.heap[parent_idx] = self.heap[parent_idx], self.heap[idx]
            idx = parent_idx

    def pop(self):
        if len(self.heap) <= 1:
            print(0)
        else:
            print(self.heap[1])

            self.heap[1] = self.heap[-1]
            del self.heap[-1]

            idx = 1
            while self.move_down(idx):
                left_child_idx = idx * 2
                right_child_idx = idx * 2 + 1

                # Case2: 오른쪽 자식만 없을 때
                if right_child_idx >= len(self.heap):
                    if self.heap[idx] > self.heap[left_child_idx]:
                        self.heap[idx], self.heap[left_child_idx] = self.heap[left_child_idx], self.heap[idx]
                        idx = left_child_idx

                # Case3: 왼쪽, 오른쪽 자식 모두 있을 때
                else:
                    if self.heap[left_child_idx] < self.heap[right_child_idx]:
                        if self.heap[idx] > self.heap[left_child_idx]:
                            self.heap[idx], self.heap[left_child_idx] = self.heap[left_child_idx], self.heap[idx]
                            idx = left_child_idx
                    else:
                        if self.heap[idx] > self.heap[right_child_idx]:
                            self.heap[idx], self.heap[right_child_idx] = self.heap[right_child_idx], self.heap[idx]
                            idx = right_child_idx


h = Heap()
N = int(stdin.readline())
for _ in range(N):
    num = int(stdin.readline())
    if num == 0:
        h.pop()
    else:
        h.insert(num)

 

마지막으로 자바에서 직접 힙 구조를 구현한 코드는 다음과 같습니다.

import java.io.*;
import java.util.*;

class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        MinHeap minHeap = new MinHeap();

        int N = Integer.parseInt(br.readLine());
        for(int i = 0; i < N; i++) {
            int num = Integer.parseInt(br.readLine());

            if (num == 0) {
                minHeap.pop();
            } else {
                minHeap.insert(num);
            }
        }
    }
}

class MinHeap {
    List heap;

    MinHeap() {
        heap = new ArrayList<>();
        heap.add(null);
    }

    boolean move_up(int idx) {
        if (idx <= 1) {
            return false;
        }

        int parent_idx = idx / 2;
        if ((int) heap.get(idx) < (int) heap.get(parent_idx)) {
            return true;
        } else {
            return false;
        }
    }

    boolean move_down(int idx) {
        int left_child_idx = idx * 2;
        int right_child_idx = idx * 2 + 1;

        // Case1: 왼쪽 자식도 없을 때
        if (left_child_idx >= heap.size()) {
            return false;
        } else if (right_child_idx >= heap.size()) { // Case2: 오른쪽 자식만 없을 때
            if ((int) heap.get(idx) > (int) heap.get(left_child_idx)) {
                return true;
            } else {
                return false;
            }
        } else { // Case3: 왼쪽, 오른쪽 자식 모두 있을 때
            if ((int) heap.get(left_child_idx) < (int) heap.get(right_child_idx)) {
                if ((int) heap.get(idx) > (int) heap.get(left_child_idx)) {
                    return true;
                } else {
                    return false;
                }
            } else {
                if ((int) heap.get(idx) > (int) heap.get(right_child_idx)) {
                    return true;
                } else {
                    return false;
                }
            }
        }
    }

    void insert(int data) {
        heap.add(data);

        int idx = heap.size() - 1;
        while (move_up(idx)) {
            int parent_idx = idx / 2;
            int tmp = (int) heap.get(idx);
            heap.set(idx, heap.get(parent_idx));
            heap.set(parent_idx, tmp);
            idx = parent_idx;
        }
    }

    void pop() {
        if (heap.size() <= 1) {
            System.out.println(0);
        } else {
            System.out.println(heap.get(1));

            heap.set(1, heap.get(heap.size() - 1));
            heap.remove(heap.size() - 1);

            int idx = 1;
            while (move_down(idx)) {
                int left_child_idx = idx * 2;
                int right_child_idx = idx * 2 + 1;

                if (right_child_idx >= heap.size()) { // Case2: 오른쪽 자식만 없을 때
                    if ((int) heap.get(idx) > (int) heap.get(left_child_idx)) {
                        int tmp = (int) heap.get(idx);
                        heap.set(idx, heap.get(left_child_idx));
                        heap.set(left_child_idx, tmp);
                        idx = left_child_idx;
                    }
                } else { // Case3: 왼쪽, 오른쪽 자식 모두 있을 때
                    if ((int) heap.get(left_child_idx) < (int) heap.get(right_child_idx)) {
                        if ((int) heap.get(idx) > (int) heap.get(left_child_idx)) {
                            int tmp = (int) heap.get(idx);
                            heap.set(idx, heap.get(left_child_idx));
                            heap.set(left_child_idx, tmp);
                            idx = left_child_idx;
                        }
                    } else {
                        if ((int) heap.get(idx) > (int) heap.get(right_child_idx)) {
                            int tmp = (int) heap.get(idx);
                            heap.set(idx, heap.get(right_child_idx));
                            heap.set(right_child_idx, tmp);
                            idx = right_child_idx;
                        }
                    }
                }
            }
        }
    }
}