[알고리즘] 백준 11279 최대 힙 (파이썬 풀이)

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

11279번: 최대 힙

문제

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

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

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

입력

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

출력

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

 

해결코드

import sys
import heapq
input = sys.stdin.readline

heap = []
n = int(input())
for i in range(n):
    x = int(input())
    if x==0:
        if heap==[]:
            print(0)
        else:
            print(heapq.heappop(heap)[1])
    else:
        heapq.heappush(heap, (-x, x))

문제풀이

이전에 풀었던 최소 힙과 반대되는 문제인데, 조금 더 난이도가 있다.

heapq 라이브러리를 사용하면 heappop을 사용했을 때 무조건 리스트 중 최소값만 제거되므로

최대값을 제거하기 위해서는 heapq라이브러리를 응용 시켜야 한다.

 

코드를 한 줄씩 보면서 설명하면

heap 이라는 리스트를 하나 생성해준다.

그리고 변수 n에 실행횟수를 입력받아준다.

n만큼 for문을 실행시키는데

여기서 이제 heap에 들어갈 수 x를 입력받아준다.

먼저 else문을 보면

        heapq.heappush(heap, (-x, x)) <-- 이렇게 작성했는데

이 코드의 원리는 heap에 x를 (-x, x)의 형태로 삽입하게 되면

heapq라이브러리의 특성에 따라 만약 x가 최대값이면 -x는 최소값을 띄게 되므로

-x가 최소인 값이 리스트의 인덱스 맨 앞에 오게 된다.

 

ex) x를 순서대로 1,2,3,4,5 입력했다면 -x는 -1,-2,-3,-4,-5 가 되는데

그렇게 되면 -5가 최소 값이 되므로 리스트 맨 앞에 위치하게 된다는 의미이다.

--> heap = [(-5, 5), (-4, 4), (-2, 2), (-1, 1), (-3, 3)]

 

만약 x 값이 0이고 heap이 비어있다면

0을 출력하면 된다.

만약 heap이 비어있지 않다면

heapq.heappop(heap)[1] 을 통해 heap의 최소값중 1번째 인덱스 값을 출력하게 된다.

 

ex) heap = [(-5, 5), (-4, 4), (-2, 2), (-1, 1), (-3, 3)] 이고 

print(heapq.heappop(heap)[1]) 을 실행하게 되면

여기서 최소값은 (-5,5)인데 1번째 인덱스 값은 5이기 때문에 5를 출력하고 제거하게 됨.

 

이 원리를 통해 n만큼 for문을 실행시키면 문제 해결이다.