[python] 1927: 최소 힙

문제와 조건은 아래와 같다.

0. 연산의 개수 N 입력
1. x가 자연수라면, 배열에 x라는 값을 추가
2. x가 0이라면, 배열에서 가장 값을 출력, 그 값을 제거
3. 배열이 비어있는 경우라면, 가장 작은 값을 출력하라고 한 경우(2번이라면), 0을 출력

이렇게 조건을 생각했고, 처음에는 아래와 같은 코드로 짰다.

n = int(print)
a = []

for _ in range(n):
    x = int(print)
    if len(a)==0 :
        if x ==0:
            print('0')
        else:
            a.append(x)
    elif x == 0 :
        print(min(a))
        a.remove(min(a))
    elif x > 0:
        a.append(x)

하지만 시간초과가 떴구, 시간 최적화를 위해 항상 sys.stdin.readline() 을 사용해야겠다고 다짐을 합니다.

import sys

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

for _ in range(n):
    x = int(sys.stdin.readline())
    if len(a)==0 :
        if x ==0:
            print('0')
        else:
            a.append(x)
    elif x == 0 :
        print(min(a))
        a.remove(min(a))
    elif x > 0:
        a.append(x)

아차 싶어서 이렇게 고쳤으나? 역시나 안됨.
문제는 min() 함수였다.
최소 힙 이라는 문제에 걸맞게 heapq를 사용해야 했다.
코테와 백준 초보는 문제를 많이 풀어봐야겠다고 또또 다짐합니다.

import sys
import heapq

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

for _ in range(n):
    x = int(sys.stdin.readline())
    if x ==0:
        if len(a)==0:
            print('0')
        else:
            print(heapq.heappop(a))
    else:
        heapq.heappush(a,x)

코드를 조금 다듬고, heapq 모듈을 사용하여 우선순위 큐를 구현할 수 있었다.

• heapq.heappush(a,x) : x를 리스트 a에 추가
• heapq.heappop(a) : 힙의 최솟값을 제거하고 반환.
  빈 힙에서 heappop() 을 호출하면 IndexError가 발생할 수 있음!

이렇게 구현할 경우 리스트에서 최소값을 찾고 삭제하는 과정을 O(log N) 시간에 처리할 수 있다고 합니다..

이렇게 세 번에 걸쳐 풀게 된 문제였습니다.