Koala - 5기/코딩테스트 준비 스터디

[BOJ / Python] 11286 - 절댓값 힙

IT Act. 2022. 2. 14. 03:23

문제 링크

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

 

11286번: 절댓값 힙

첫째 줄에 연산의 개수 N(1≤N≤100,000)이 주어진다. 다음 N개의 줄에는 연산에 대한 정보를 나타내는 정수 x가 주어진다. 만약 x가 0이 아니라면 배열에 x라는 값을 넣는(추가하는) 연산이고, x가 0

www.acmicpc.net


풀이

  • 파이썬의 heapq 모듈을 사용해 우선순위 큐를 구현하였다. 
  • 처음에는 양의 정수와 음의 정수를 각각 다른 힙에 저장하여 풀었다. (코드1 참조)
  • 다른 블로그를 보니 힙에 push할 때 (절댓값, 입력값) 쌍으로 저장하여 푼 풀이가 있어서 다시 풀어보았다. (코드2 참조)

코드1

import sys
input = sys.stdin.readline
from heapq import *

n = int(input())
hq1 = [] # 양의 정수
hq2 = [] # 음의 정수

for _ in range(n):
    x = int(input())
    if x>0: # 양의 정수를 입력했을 때 
        heappush(hq1, x)
    elif x<0: # 음의 정수를 입력했을 때 
        heappush(hq2, -x) # 작은 절댓값부터 정렬되어야하므로 앞에 마이너스를 붙인다. 
    elif not hq1 and not hq2: # 두 힙 모두 비어 있을 때 
        print(0)
    elif not hq1: # 양의 정수 힙이 비어있을 때 
        print(-heappop(hq2)) # 음의 정수 힙에서 마이너스를 붙이고 출력 
    elif not hq2: # 음의 정수 힙이 비어있을 때 
        print(heappop(hq1)) # 양의 정수 힙에서 출력 
    else: # 두 힙이 모두 비어있지 않을 때 
        x1 = hq1[0] # 힙에서 절대값이 가장 작은 값을 저장 
        x2 = hq2[0]
        if x1 >= x2: # 두 힙에서 나온 값을 비교. 
            # 절댓값이 같거나 양의정수 절대값이 큰 경우 가장 작은 수를 출력해야하므로 음의 정수 힙에서 pop 
            print(-heappop(hq2))
        else: # 아닐 경우 양의 정수 힙에서 pop
            print(heappop(hq1))

코드2

import sys
input = sys.stdin.readline
from heapq import *

n = int(input())
hq = []

for _ in range(n):
    x = int(input())
    if x!= 0: # 0이 아닌 수 입력 시 힙에 (절댓값, 입력값)쌍 저장 
        heappush(hq, (abs(x), x))
    elif not hq: # 힙이 비어있을 때 0 출력 
        print(0)
    else: # 비어 있지 않을 때 (절댓값, 입력값) 중 입력값 출력 
        print(heappop(hq)[1])