백준 1202 보석 도둑 (Python)

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

1202번: 보석 도둑

당연히 heap을 사용해야 한다는 것은 알고 있었지만, 어떻게 사용할 지에 대해서 고민을 많이 했었다.

보석과 가방을 같은 heap에 넣는다는 생각을 하는 게 좀 힘들었던 것 같다.

 

무게를 기준으로 최소 heap에 보석과 가방을 넣어준다.

가방 같은 경우는 같은 무게에서 가장 뒤로 갈 수 있도록 가치에는 무한을 넣어준다.

 

하나씩 꺼낼 때 쓸 heap도 만들어준다.

 

보석과 가방을 넣은 heap에서 하나씩 꺼내어서 반복문을 실행해준다.

 

만약 꺼낸 것이 보석이라면 최대 heap에 가치를 넣어준다.

만약 꺼낸 것이 가방이라면 최대 heap에서 하나의 보석을 넣을 수 있는 것이므로 가장 큰 값을 가져와 결과값에 더해준다.

 

그렇게해서 구한 결과값을 출력해주면 된다.

import sys
import heapq

N, K = map(int, sys.stdin.readline().split())

jewelry = list()

for _ in range(N):
    heapq.heappush(jewelry, list(map(int, sys.stdin.readline().split())))

for _ in range(K):
    heapq.heappush(jewelry, [int(sys.stdin.readline().strip()), float('inf')])

result = 0
cur_jewelry = list()

while jewelry:

    weight, value = heapq.heappop(jewelry)

    if value == float('inf'):
        if cur_jewelry:
            result += heapq.heappop(cur_jewelry)
    else:
        heapq.heappush(cur_jewelry, value * -1)


print(result * -1)