[백준11279번] 최대 힙 / Python3

문제

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

1. 배열에 자연수 x를 넣는다.
2. 배열에서 가장 큰 값을 출력하고, 그 값을 배열에서 제거한다.

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

입력

첫째 줄에 연산의 개수 N(1≤N≤100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 자연수라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0이라면 배열에서 가장 큰 값을 출력하고 그 값을 배열에서 제거하는 경우이다. 입력되는 자연수는 2^31보다 작다.

출력

입력에서 0이 주어진 회수만큼 답을 출력한다. 만약 배열이 비어 있는 경우인데 가장 큰 값을 출력하라고 한 경우에는 0을 출력하면 된다.

예제 입력 1

13
0
1
2
0
0
3
2
1
0
0
0
0
0

예제 출력 1

0
2
1
3
2
1
0
0

나의 풀이

최대 힙을 구현하는 문제로 힙에 대한 개념이 있어야 풀 수 있는 문제이다. 물론 기본으로 제공되는 힙의 모듈을 사용하면 몰라도 풀 수는 있겠지만...ㅋㅋ 여튼 최대 힙이란 부모노드의 값이 자식노드의 값보다 큰 자료구조이다.

그럼 힙에 push, pop하는 과정을 간다하게 설명해보겠다. 참고로 최대 힙을 배열을 이용하여 구현하였다.

• push

  최대 힙에서 push의 경우 넣으려는 값을 우선 마지막 노드에 넣어주고 부모의 값들과 비교를 하여 자식의 노드가 부모의 노드의 값보다 크지않을때까지 바꾸어 올라가면 된다. 아래 그림을 보면 이해가 더 쉬울 것이다.

  

• pop

  pop은 마지막 노드의 값을 1번 노드에 넣어주고 push와 반대로 자식들과 비교하면서 바꾸어 내려가는 것이다. 여기서 주의할 점이 있다면, 한 부모의 자식이 둘일 경우 왼쪽과 오른쪽 자식 중에 더 큰 자식과 바꾸어 내려가야 한다는 것이다.

  마찬가지로 그림을 첨부하겠다.

  

이렇게 최대 힙에 대한 설명을 마치고, 아래 코드를 보여주겠다. 그런데 검색을 해보니 python의 heapq를 사용하면 간편하게 이 문제를 풀 수 있었다. heapq는 최소힙으로 구현이 되어있어 이 문제를 풀기위해선 입력받는 값에 음수 값을 취해서 push를 해주고 pop을 해주면 풀 수 있었다. 이 코드도 아래 첨부하겠다.

아! 그리고 input()을 사용하면 시간초과가 날 수도 있다... 필자도 input()을 이 문제를 푸니 시간초과가 나와 혹시나 sys.stdin.readline()를 사용하니 시간초과가 나지 않았다.

코드

• 최대 힙 직접 구현

# 11279번 최대 힙
from math import *
import sys

# insert
def insert(heap, num):
    heap.append(num)

    i = len(heap) - 1
    while i > 1:
        if heap[i] > heap[i//2]:
            tmp = heap[i]
            heap[i] = heap[i//2]
            heap[i//2] = tmp
            i = i // 2
        else:
            break

# remove
def remove(heap):
    maxVal = heap[1]
    tmp = heap.pop()

    parent = 1
    child = 2

    while child <= len(heap)-1:
        if child < len(heap)-1 and heap[child] < heap[child+1]:
            child += 1

        if heap[child] <= tmp:
            break

        heap[parent] = heap[child]

        parent = child
        child *= 2

    if len(heap) != 1:
        heap[parent] = tmp
    return maxVal

# main
n = int(sys.stdin.readline())

heap = [0]

for _ in range(n):
    num = int(sys.stdin.readline())

    if num == 0:
        if len(heap) == 1:
            print(0)
        else:
            print(remove(heap))
    else:
        insert(heap, num)

• heapq 사용

# 11279번 최대 힘 (heapq사용)
import sys
import heapq

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

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