[백준] 최대 힙(python)

백준, 최대 힙

11279번: 최대 힙

TL;DR

• 우선순위 큐(Priority queue)

문제 요약

1. 입력받은 X가 자연수라면 배열에 넣고 0이라면 배열에 있는 값 중 가장 큰 값을 출력하고 값을 배열에서 제거한다.

 

- 입력에 따라 수행되는 구문이 다르며, 0이 입력되었을 때 배열 내에서 가장 큰 수를 출력해야 한다.

- 이를 완전 탐색으로 풀 때, 최악의 경우 O(n^2)가 되므로 다른 방법을 찾아야 한다.

- 가장 큰 값에 우선순위를 두어 우선순위 큐를 사용하면 문제를 해결할 수 있다.

입출력 형태

입출력 예시

풀이

파이썬에서 우선순위 큐는 heapq로 편하게 사용할 수 있다.

힙(heap)은 트리 기반의 자료구조로 최소 힙(min-heap)과 최대 힙(max-heap)으로 나눌 수 있다. 이진 탐색 트리와는 다르게 노드들의 정렬된 형태가 아니며, 규칙에 따라 노드의 삽입-삭제가 이루어지는 자료구조형이다.

최소 힙의 경우는 루트가 힙 내에서 가장 작은 값을 가지며, 부모 노드는 자식 노드에 비해서 작은 값을 갖는다. 최대 힙은 그 반대로 루트가 힙 내에서 가장 큰 값을 가지며, 부모 노드는 자식 노드에 비해 큰 값을 갖는다.

파이썬에서의 힙은 heapq로 쉽게 사용할 수 있다. 하지만 기본적으로 최소 힙(min heap)으로 구현되어 있기 때문에 heappop을 할 경우 가장 작은 값을 빼내게 된다. 문제 해결을 위해서는 입력을 바꾸는 방식으로 문제에 접근해야 한다.

import sys, heapq

N = int(sys.stdin.readline())
heap = []

for _ in range(N):
    X = int(sys.stdin.readline())
    if X == 0:
        try:
            print(-heapq.heappop(heap))
        except:
            print(0)
    else:
        heapq.heappush(heap, -X)

 

이를 막기 위해 코드에서는 입력을 할 때 -를 붙여서 절대값이 큰 수가 작은 수가 되도록 만든다. 이렇게 하면 heappop을 했을 때 절대값이 큰 수가 높은 우선순위로 배열에서 빠져나오게 할 수 있다. 출력을 할 때는 다시 -기호를 붙여 원래의 숫자로 출력하도록 작성해주면 된다.

 

try-except 구문은 비어있는 배열에 대해 pop을 수행할 경우 에러가 발생하기 때문에 사용하였다. 다음과 같은 구문으로 대체할 수 있다.

if len(heap) == 0:
    print(0)
else:
    print(-heapq.heappoop(heap))