백준 11286번 절대값 힙 [Python]

[Silver I] 절댓값 힙 - 11286

문제 링크

성능 요약

메모리: 39824 KB, 시간: 140 ms

분류

자료 구조, 우선순위 큐

제출 일자

2024년 3월 20일 22:04:02

문제 설명

절댓값 힙은 다음과 같은 연산을 지원하는 자료구조이다.

1. 배열에 정수 x (x ≠ 0)를 넣는다.
2. 배열에서 절댓값이 가장 작은 값을 출력하고, 그 값을 배열에서 제거한다. 절댓값이 가장 작은 값이 여러개일 때는, 가장 작은 수를 출력하고, 그 값을 배열에서 제거한다.

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

입력

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

출력

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

코드

import sys, heapq
input = sys.stdin.readline
print = sys.stdout.write

arr = []
N = int(input())
for _ in range(N):
    x = int(input())
    if x:
        heapq.heappush(arr, (abs(x), x))
    else:
        if arr:
            print("%d\n"%heapq.heappop(arr)[1])
        else:
            print("%d\n"%0)

처음에 heapq를 사용하지 않고 리스트로 코드를 구성하다가 시간초과가 일어났다. 다른 분들의 풀이를 참고해서 파이썬에서 지원하는 heapq를 이용해서 코드를 구성했다.

N횟수 만큼 x값을 입력받고, x가 0이 아니면 heapq의 heappush연산을 이용한다. 두번째 인자에 튜플형태로 (abs(x), x) 넣어 줌으로써 절대값 기준으로 정렬될 수 있게 했다. x를 절대값 연산한 값과 원래 값을 비교하여 값이 작은 것을 arr리스트에 추가한다. x가 0인 경우에는 arr리스트가 비어있는지 확인하고, 비어있지 않다면, heappop연산으로 리스트 내에 있는 가장 작은 값을 제거하는데, 이때 튜플 형태로 저장되어 있으므로 절대값을 씌우지 않은 [1]번째 값을 출력한다. arr리스트가 비어있는 경우에는 출력할 원소가 없으므로 0을 출력해준다.