[백준 알고리즘][파이썬/Python] 11279번: 최대 힙

문제

시간 제한 : 1 초 (추가 시간 없음), 메모리 제한 : 256 MB

 

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

 

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

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

 

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

 

입력

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

 

출력

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

 

입출력 예

 

제출 코드

import sys
import heapq as hq

n = int(input())
heap = []
for i in range(n):
    x = int(sys.stdin.readline())
    if x:
        hq.heappush(heap, (-x, x))
    else:
        if len(heap) >= 1:
            print(hq.heappop(heap)[1])
        else:
            print(0)

 

문제의 핵심

# 1

Python의 내장 모듈인 heapq를 활용하여 효율적인 heap 연산을 할 수 있어야 한다. 물론 heap 자료 구조에 대한 개념부터 알아야 하는 것이 먼저다.

 

# 2

heapq 모듈에서는 최소 힙(Min Heap)만을 지원하는데, 문제에서 요구하는 것을 시간 안에 해결하기 위해서는 최대 힙(Max Heap)을 구현해야 한다. 가장 효율적으로 최대 힙을 구현할 수 있어야 하는데, 여기서 계속 버벅였다. 생각을 하다가 레퍼런스를 참고하여 문제를 해결했는데 힌트를 얻은 곳은 [여기]다.

 

 

 

라인 별 풀이

(1) line 7

시간 제한이 엄격하게 1초로 정해져 있는데 최대 입력 개수가 10만개이므로 sys 모듈을 활용하여 입력을 받았다.

 

(2) line 8, 9

이 두 줄이 해당 문제에 핵심이다. 입력값이 0이 아니면 해당 값을 heap에 push해야 한다. heapq 모듈에 정의되어 있는 heappush(heap, value)는 heap에 value를 push하면서 heap의 성질을 유지시키는 역할을 해준다. 하지만 value만을 push하게 될 경우 최소 힙의 성질을 갖게 된다. 즉, 가장 작은 값이 heap의 루트 노드로 오게끔 되어 있다.

 

여기서 (-value, value)로 구성된 튜플을 만들어서 heappush 함수의 인자로 주게 되면 튜플의 0번째 원소인 -value를 기준으로 heap을 구성하게 된다. 기존의 value가 가장 클수록 -value의 값이 가장 작은 값이 되어 루트 노드에 올 수 있게 된다.

 

ex) input으로 3, 1, 5이 주어지면 아래와 같은 순서대로 힙을 구성하게 된다.

heap = [(-3, 3)]

heap = [(-3, 3), (-1, 1)]

heap = [(-5, 5), (-1, 1), (-3, 3)]

 

그 덕분에 필요할 때마다 heapify를 하지 않아도 되며 결과적으로 소스 코드의 전체적인 시간 복잡도도 훨씬 작아지게 된다.

 

(3) line 10 ~ 14

입력값이 0일 경우, heap의 크기가 1 이상이면 heap에서 가장 큰 값을 출력하며 pop시켜야 한다. heap에 아무런 값도 들어있지 않다면 0을 출력하면 된다. 여기서 중요한 것은 가장 큰 값을 출력할 때 hq.heappop(heap)만 쓰면 안 된다는 것이다. hq.heappop(heap)은 튜플을 반환하기 때문에 해당 튜플에서 본래 입력된 값을 출력해줘야 한다. 그래서 두 번째 요소를 반환해야 한다.

 

 

 

실패 기록

import sys
import heapq as hq

n = int(input())
heap = []
for i in range(n):
    x = int(sys.stdin.readline())
    if x:
        heap.append(x)
    else:
        if len(heap) >= 1:
            hq._heapify_max(heap)
            print(hq.heappop(heap))
        else:
            print(0)

heapq 모듈 안에 최대 힙으로 heapify해주는 함수가 있다. 실패했던 코드에서는 필요할 때마다 _heapify_max 함수를 썼는데, heapify의 시간 복잡도는 O(N)이므로 1초라는 제한 시간 안에서 해결하기 위해 heapify를 최대한 지양해야 했다.

 

 

 

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