[Python] 11279 - 최대 힙

문제

11279번: 최대 힙

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

널리 잘 알려진 자료구조 중 최대 힙이 있다. 최대 힙을 이용하여 다음과 같은 연산을 지원하는 프로그램을 작성하시오.

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

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

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

입력

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

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

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

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

출력

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

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

---

풀이

문제 이름에 나와있듯이 최대힙을 구현하는 문제다.

최소힙, 최대힙은 클래스로 직접 구현할수도 있고 python의 heapq모듈을 사용할수도 있다.

최소힙포스팅에서 클래스로 구현하는 방법을 시도해봤으므로 최대힙은 heapq모듈을 사용해 풀어봤다.

설명

• 입력값이 0이 아니면 힙에 값을 넣는데, heapq모듈을 최소힙만 제공한다. 따라서 튜플로 값을 묶어 부호를 바꿔준 값과 함께 저장한다.

• 이렇게 저장하는 이유는 heapq모듈이 단일 값이 아닌 묶음이 들어올 때 제일 앞에 있는 값을 기준으로 순서를 정하기때문이다. (튜플의 0번째 값을 역수로 주면 최소→최대로 변경)

• 입력되는 숫자가 0이면 힙에서 제거해야한다.
  • 힙이 비었으면 0을 저장하고
  • 아니면 제일 큰 값을 뺀다.

답

from sys import stdin
import heapq

input = stdin.readline


def max_heap(x):
    h = []
    ret = []
    for num in x:
        if num == 0:
            ret.append(0 if not h else heapq.heappop(h)[1])
        else:
            heapq.heappush(h, (-num, num))
    return ret


if __name__ == "__main__":
    N = int(input())
    X = []
    for _ in range(N):
        X.append(int(input()))
    res = max_heap(X)
    print('\n'.join(map(str, res)))

Uploaded by Notion2Tistory v1.1.0