[Python] 1927 - 최소 힙

문제

1927번: 최소 힙

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

널리 잘 알려진 자료구조 중 최소 힙이 있다.

최소 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.

1. 배열에 자연수 x를 넣는다.

2. 배열에서 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다.

프로그램은 처음에 비어있는 배열에서 시작하게 된다.

입력

• 첫째 줄에 연산의 개수 N(1 ≤ N ≤ 100,000)이 주어진다.

• 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다.

• 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 작은 값을 출력하고 그 값을 배열에서 제거하는 경우이다.

• 입력되는 자연수는 231보다 작다.

출력

• 입력에서 0이 주어진 회수만큼 답을 출력한다.

• 만약 배열이 비어 있는 경우인데 가장 작은 값을 출력하라고 한 경우에는 0을 출력하면 된다.

---

풀이

• 최소 힙에 대한 정직한 문제다.

• 최소 힙을 구현해본적이 없다면 priority queue모듈을 사용해서 풀기전에 직접 구현해서 문제를 풀어보길 바란다.

• 최소 힙과 완전이진트리에 대한 개념은 많으므로 해당 개념을 먼저 숙지하고 아래 코드 설명을 봐야 이해가 가도록 작성하였다.

최소 힙 구현

• 최소 힙을 클래스로 정의하고 insert, delete, swap, sort 함수를 구현한다.

• 최소 힙은 완전이진트리를 배열형태로 구현하는 것이므로 배열의 인덱스를 2의 배수로 활용하면 배열을 트리처럼 사용할 수 있다.

  완전이진트리는 2개의 자식만을 가진다.

  • 따라서 부모의 인덱스는 i // 2
  • 왼쪽 자식의 인덱스는 i * 2
  • 오른쪽 자식의 인덱스는 i * 2 + 1

  이 된다.

• insert 함수
  1. 맨 마지막에 원소 추가
  2. 부모보다 클때까지 이동하면서 swap

• delete 함수
  1. 배열이 비었으면 0반환
  2. 아니라면 맨 처음이랑 마지막 원소 swap
  3. 맨 처음이었던 원소가 swap되면서 마지막으로 이동했으니 마지막 원소를 pop
  4. 맨 마지막이었던 원소가 맨 처음으로 갔으니 배열을 재정렬

• swap 함수

  원소의 인덱스를 받아 서로 교환

• sort 함수
  1. 첫 인덱스 1부터 자식과 비교
  2. 왼쪽 자식과 비교해서 왼쪽 자식이 더 작다면 바꿀 인덱스 _next를 왼쪽 자식의 인덱스로
  3. 오른쪽 자식도 왼쪽 자식과 같은 로직 수행
  4. 바꿀 인덱스 _next가 현재 인덱스 _idx랑 다르다면 swap
  5. 다르지 않다면 더 이상 바꿀게 없으므로 break

답

from sys import stdin

input = stdin.readline


class MinHeap:

    def __init__(self):
        self.que = [0]

    def insert(self, x):
        # 맨 마지막에 원소 추가
        self.que.append(x)
        # 부모보다 클때까지 이동
        idx = len(self.que) - 1  # 현재 인덱스
        while 1 < idx:
            # 부모와 비교 (부모 : 현재 // 2)
            if self.que[idx] < self.que[idx // 2]:  # 부모가 더 크면 swap
                self.swap(idx, idx // 2)
                idx = idx // 2
            else:  # 부모가 더 작으면 끝
                break

    def delete(self):
        # 배열이 비었을때
        if len(self.que) == 1:
            return 0
        # 맨 처음이랑 맨 마지막 swap
        self.swap(1, -1)
        # 맨 마지막이 최소이므로 pop
        _min = self.que.pop()
        # 재정렬
        self.sort()
        return _min

    def swap(self, i1, i2):
        self.que[i1], self.que[i2] = self.que[i2], self.que[i1]

    def sort(self):
        _idx = 1

        while len(self.que) > 1:
            _next = _idx
            if _idx * 2 < len(self.que) and self.que[_idx * 2] < self.que[_next]:
                _next = _idx * 2

            if _idx * 2 + 1 < len(self.que) and self.que[_idx * 2 + 1] < self.que[_next]:
                _next = _idx * 2 + 1

            if _next != _idx :
                self.swap(_idx, _next)
                _idx = _next

            else:
                break


if __name__ == "__main__":
    N = int(input())
    pq = MinHeap()
    for _ in range(N):
        x = int(input())
        if x == 0:
            print(pq.delete())
        else:
            pq.insert(x)

Uploaded by Notion2Tistory v1.1.0