백준 1202 파이썬 - 보석 도둑

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

1202번: 보석 도둑

문제

세계적인 도둑 상덕이는 보석점을 털기로 결심했다.

상덕이가 털 보석점에는 보석이 총 N개 있다. 각 보석은 무게 Mi와 가격 Vi를 가지고 있다. 상덕이는 가방을 K개 가지고 있고, 각 가방에 담을 수 있는 최대 무게는 Ci이다. 가방에는 최대 한 개의 보석만 넣을 수 있다.

상덕이가 훔칠 수 있는 보석의 최대 가격을 구하는 프로그램을 작성하시오.

코드

import sys
import heapq
n, k = map(int, sys.stdin.readline().split())
mv = [list(map(int, sys.stdin.readline().split())) for _ in range(n)]
k = [int(sys.stdin.readline()) for _ in range(k)]
mv.sort()
k.sort()
ans = 0
q = []
for i in k:
    while mv and mv[0][0] <= i:
        heapq.heappush(q, -heapq.heappop(mv)[1])
    if q:
        ans -= heapq.heappop(q)
sys.stdout.write(str(ans))

풀이

이 문제는 우선순위큐를 사용해 풀어야 시간초과가 나지 않을 가능성이 높다고 한다. 우선 heapq를 import해주고, q에다가 무게가 가장 적은 보석의 가격을 넣는다. pop을 하면 가장 작은 값을 우선적으로 가져올 수 있다.(파이썬에서의 heapq는 min heap)  가격 최대값을 구해야 하기 때문에 마이너스를 붙여 넣는다. 나중에 pop을 하면 가장 작은 숫자는 절대값이 가장 큰 음수가 될 것이다. 그리고 그 수를 ans에 빼주어 양수가 되게 한다.