[BOJ 1927] 최소 힙

문제 요약

문제 요약
최소 힙을 사용하여 다음 연산을 만족하는 프로그램을 작성하라
1 배열에 자연수를 넣는다
2 자연수 대신 0을 입력하면 배열에 가장 작은 값을 출력하고 제거한다
(만약 배열이 비었을 때 0을 입력하면 0을 출력한다)
입력
첫째 줄 : 입력할 수의 개수
둘째 줄 ~ N+1째 줄: 입력할 수들
출력
0 입력할 때 나오는 출력들

풀이

'배열에 자연수를 넣는다', '가장 작은 값을 출력하고 제거한다'가 가능한 최소 힙을 구현한다

최소 힙은 어떻게 구현하면 될까?

최소 힙은 완전 이진 트리를 통해 구현한다.
하지만 파이썬에서는 heapq라는 내장 모듈이 있기 때문에 직접 구현할 필요는 없다.
최소 힙에 자연수를 넣는 것은 heappush라는 메서드로, 제거하는 것은 heappop이라는 메서드로 구현할 수 있다.

from heapq import heappush, heappop
heap = []
heappush(heap, 1)
print(heappop(heap))

어떤 식으로 구현하면 될까?

1. 자연수를 입력받는다
   1-1. heappush를 통해 넣는다
2. 0을 입력받는다
   2-1. 배열이 안 비어있으면 출력하고 제거한다
   2-2. 배열이 비어있으면 0을 출력한다

신경써야 할 부분

입력수가 10만이므로 sys.stdin.readline을 사용한다
heappop은 배열이 비었을 때 꺼내려고 하면 IndexError를 띄우므로 try except를 통해 해결한다

결과 코드

import sys
from heapq import heappush, heappop
input = sys.stdin.readline

N = int(input())
max_heap = []

for i in range(N):
    num = int(input())
    if num != 0:
        heappush(max_heap, num)
    else:
        try:
            print(heappop(max_heap))
        except IndexError:
            print(0)

---

[BOJ 11279] 최대 힙

이 문제는 최소 힙과 쌍둥이 문제로 heappop을 쓸 때 최솟값이 아닌 최댓값을 출력하도록 하면 된다

heapq로 최대 힙을 어떻게 구현할까?

heapq는 기본적으로 최소 힙으로 구현된다
heappush를 할 때 num 대신 (-num, num)이라는 튜플을 집어넣으면 튜플의 첫 원소 기준으로 정렬하게 된다
부호가 반대이므로 자연스레 최대 힙과 같이 정렬되게 된다
그리고 heappop으로 꺼낼 때는 heappop(heap)[1]을 통해 양수를 꺼내면 된다.

결과 코드 (최대 힙)

import sys
from heapq import heappush, heappop
input = sys.stdin.readline

N = int(input())
max_heap = []

for i in range(N):
    num = int(input())
    if num != 0:
        heappush(max_heap, (-num, num))
    else:
        try:
            print(heappop(max_heap)[1])
        except IndexError:
            print(0)