99클럽 코테 스터디 파이썬 비기너 18일차 크리스마스 선물

문제

문제
크리스마스에는 산타가 착한 아이들에게 선물을 나눠준다. 올해도 산타는 선물을 나눠주기 위해 많은 노력을 하고 있는데, 전세계를 돌아댕기며 착한 아이들에게 선물을 나눠줄 것이다. 하지만 산타의 썰매는 그렇게 크지 않기 때문에, 세계 곳곳에 거점들을 세워 그 곳을 방문하며 선물을 충전해 나갈 것이다. 또한, 착한 아이들을 만날 때마다 자신이 들고있는 가장 가치가 큰 선물 하나를 선물해 줄 것이다.

이제 산타가 선물을 나눠줄 것이다. 차례대로 방문한 아이들과 거점지의 정보들이 주어졌을 때, 아이들이 준 선물들의 가치들을 출력하시오. 만약 아이들에게 줄 선물이 없다면 -1을 출력하시오.

입력
첫 번째 줄에서는 아이들과 거점지를 방문한 횟수 n이 주어진다.(1≤n≤5,000)

다음 n줄에는 a가 들어오고, 그 다음 a개의 숫자가 들어온다. 이는 거점지에서 a개의 선물을 충전하는 것이고, 그 숫자들이 선물의 가치이다. 만약 a가 0이라면 거점지가 아닌 아이들을 만난 것이다. 선물의 가치는 100,000보다 작은 양의 정수이다.(1≤a≤100)

출력
a가 0일 때마다, 아이들에게 준 선물의 가치를 출력하시오. 만약 줄 선물이 없다면 -1을 출력하라. 적어도 하나의 출력이 있음을 보장한다.

예제 입력 1

5 0 2 3 2 0 0 0

 

예제 출력 1

-1 3 2 -1

알고리즘 분류
자료 구조, 우선순위 큐

From Baekjoon, https://www.acmicpc.net/problem/14235

답안

import sys
import heapq

gifts = [] ## Integer
results = [] ## String
cnt = int(sys.stdin.readline())

for _ in range(cnt):
    input_data = sys.stdin.readline().strip() ## Name 'input' X

    if input_data == '0':
        result = -heapq.heappop(gifts) if gifts else "-1" ## Exception ## - to +
        results.append(str(result))

    else: ## Max heap required
       replenishers = input_data.split()
       replenishers.pop(0) ## Remove a
       for r in replenishers:
           heapq.heappush(gifts, -int(r)) ## Make an array a heap

sys.stdout.write('\n'.join(results))

원리

1. 입력받은 값이 '0'이면 힙에서 가장 큰 값을 꺼내서 준다.
2. 입력받은 값이 '0'이 아니라면 각각의 값을 음수로 바꾸어 힙에 저장해 정렬해 준다. (최대 힙 구현)

구현

• '0'이 아닌 입력 값에서 가장 첫 번째 값은 필요가 없어, pop(0)으로 제거했다. remove(replenishers[0])보다 효율적임.
• 저장할 때 값을 음수로 바꾸고, heappop(gifts)로 가져올 때 앞에 -를 붙여주면 최대 힙 구현이 더 간단하다.

오류

• 파이썬에서 'input'이라는 이름의 변수를 사용하면 input()랑 충돌의 가능성이 있으므로 input_data 등으로 대체하자.
• heapq.nlargest(K, myHeap), heapq.nsmallest(K, myHeap)를 사용할 때 주의할 점으로, 둘은 함수를 팝 하는 게 아니라서 요소는 그대로 힙에 남아있으며, 반환 값이 리스트이기 때문에 요소가 하나라 해도 result[0]처럼 사용해야 한다.

후기

오늘은 큰 문제 없이 빨리 해결해서 기분이 좋다. 복잡도를 계산하는 것도 의식적으로 연습하고 있다. 쉬운 문제를 풀고 있지만, 사고와 습관이 개선된 게 느껴진다. 다음 달에는 미들러 레벨로 넘어가야지.