[PS] BOJ 1202 / 보석 도둑

[PS] BOJ 1202 / 보석 도둑
문제 링크: https://www.acmicpc.net/problem/1202
Thumbnail: Photo by Bùi Hoàng Long (Unsplash)

정렬과 우선순위 큐를 활용해 배낭을 채우는 문제입니다.

풀이

아이디어

​​이 문제는 가방과 보석 모두 정렬한 뒤 풀어야 합니다.

가방은 가방의 용량(담을 수 있는 최대 무게)에 따라 오름차순으로,
보석은 보석의 무게에 따라 내림차순으로 정렬합니다.

gems = [] # (M, V)
N, K = map(int, input().split())

for _ in range(N):
    m, v = map(int, input().split())
    gems.append((m, v))

gems.sort(key=lambda g: g[0])
bag_weights = sorted(int(input()) for _ in range(K))

가방과 보석 정렬

이후, 각 가방마다 담을 수 있는 보석을 모두 저장하고, 그 중 최대 가치를 가진 보석을 찾습니다.

이 때, 가방을 용량에 따라 오름차순으로 정렬했기 때문에 각 가방은 이전 가방이 담을 수 있는 모든 보석을 담을 수 있음이 보장됩니다! 따라서 이전에 저장한 리스트를 재활용할 수 있습니다.

가방에 담을 수 있는 모든 보석들을 저장하는 건 우선 순위 큐를 사용했으며, 최대 가치의 보석을 먼저 꺼내기 위해 각 보석의 가치 $V$에 대해 $-V$를 우선 순위 값으로 사용했습니다. (heapq모듈은 기본적으로 크기가 작은 값을 먼저 가져옵니다.)

전체 코드

from heapq import heappush, heappop
input = open(0).readline

gems = [] # (M, V)
N, K = map(int, input().split())

for _ in range(N):
    m, v = map(int, input().split())
    gems.append((m, v))

gems.sort(key=lambda g: g[0])
bag_weights = sorted(int(input()) for _ in range(K))
bags = []
res = 0

gem_idx = 0
bag_size = 0
for bag_idx in range(K):
    while gem_idx < N and gems[gem_idx][0] <= bag_weights[bag_idx]:
        v = gems[gem_idx][1]
        heappush(bags, (-v, v)) # Sort based on gem's value. (decreasing order)
        gem_idx += 1
        bag_size += 1
    if bag_size > 0:
        _, max_gem = heappop(bags) # Last element : Max Value
        res += max_gem
        bag_size -= 1

print(res)

solution.py