백준 1202번: 보석 도둑 (G2) - Python 풀이

문제 설명

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

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

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

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

입력

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000)

다음 N개 줄에는 각 보석의 정보 Mi와 Vi가 주어진다. (0 ≤ Mi, Vi ≤ 1,000,000)

다음 K개 줄에는 가방에 담을 수 있는 최대 무게 Ci가 주어진다. (1 ≤ Ci ≤ 100,000,000)

모든 숫자는 양의 정수이다.

출력

첫째 줄에 상덕이가 훔칠 수 있는 보석 가격의 합의 최댓값을 출력한다.

예제 입력 1

2 1
5 10
100 100
11

예제 출력 1

10

예제 입력 2

3 2
1 65
5 23
2 99
10
2

예제 출력 2

164

힌트

두 번째 예제의 경우 첫 번째 보석을 두 번째 가방에, 세 번째 보석을 첫 번째 가방에 넣으면 된다.

문제 요약) 보석마다 무게와 가격이 있고, 가방에는 무게 제한이 있습니다. 하나의 가방에 하나의 보석만 담을 수 있을 때 담을 수 있는 보석의 최대 가격을 구해라.

풀이 과정

보석 무게는 적을수록 좋은데, 가격은 높을수록 좋습니다.

무게 제한 값이 작은 가방부터 하나씩 확인하는데, 그 가방에 넣을 수 있는 모든 보석들의 가격을 최대힙 구조로 저장합니다. 최대 가격을 뽑아야 하니까요.

그 가방에 넣을 수 없는 보석 A가 뜨면, 최대힙에서 보석의 최대 가격을 꺼내고 answer에 저장합니다. 이로써 현재 가방에 가능한 최대 가격의 보석을 담게 되는 겁니다.

그리고 다음 가방을 확인합니다. 앞의 가방이 넣을 수 없었던 A를 이 가방도 못 넣는다면, 이 가방도 최대힙에서 보석의 최대 가격을 꺼내고 answer에 저장합니다. 가방은 무게 제한이 작은 것부터 확인하고 있으니, 최대힙에 존재하는 모든 보석은 이 가방에도 담을 수 있는 보석들이 되기 때문이죠.

그리고 그 다음 가방을 확인하는데, 앞의 가방들이 담지 못한 A 보석을 담을 수 있는 가방이라면 A의 가격을 최대힙에 추가합니다.

이를 모든 가방을 확인할 때까지 계속 반복합니다.

최대힙에는 현재 가방에 담을 수 있는 ‘후보 보석들의 가격’이 최대힙 구조로 저장될 것입니다.

파이썬 heapq 자료구조를 사용할 때 알아야 할 점은

1. 기본적으로 최소힙으로 구성됩니다.
   • 그래서 최대힙으로 구성하려면 값에 -를 붙이면 됩니다.
2. 오름차순으로 출력되지 않습니다.
   • heapq는 리스트 전체를 정렬하는 게 아니라, 내부적으로 완전이진트리 구조에서 최소값이 루트(인덱스 0)에 있도록 유지합니다.
   • 그래서 리스트 형태로 보면 이상하게 보일 수 있지만, heap[0]이 항상 최소라는 점은 자료구조적으로 보장됩니다.

정답 코드

import sys
import heapq as hq
input = sys.stdin.readline

N, K = map(int, input().split())

# 보석 정보를 (보석 무게, 보석 가격) 튜플 형태로 저장할 리스트
# 보석 무게를 기준으로 최소힙이 될 것
jewels = []
# 가방 무게 제한을 담을 리스트
bag_limits = []

answer = 0

# M 값을 기준으로 최소힙 jewels 구성
for _ in range(N):
    M, V = map(int, input().split())
    hq.heappush(jewels, (M, V))

# 가방 무게 제한 리스트 구성하고
for _ in range(K):
    bag_limits.append(int(input().strip()))
# 오름차순 정렬
bag_limits.sort()

# 보석 가격을 최대힙 구조로 저장할 리스트 생성
jewel_values = []

# 무게 제한 작은 가방부터 확인
for limit in bag_limits:
    # 담을 보석이 아직 있고, 가장 가벼운 보석이 현재 가방에 들어가면
    while jewels and limit >= jewels[0][0]:
        # 가장 가벼운 보석을 jewels에서 빼고
        # 해당 보석의 무게값을 - 붙여서 (최대힙으로 만들기 위해)
        # jewel_values에 넣는다
        hq.heappush(jewel_values, -hq.heappop(jewels)[1])
    # 최대힙에 가격이 남아 있으면
    if jewel_values:
        # 보석의 최대 가격을 answer에 더한다
        # 원래 값으로 돌리기 위해 -를 붙인다
        answer += -hq.heappop(jewel_values)

print(answer)

가방의 무게 배열과 보석 정보 배열은 무게 값을 기준으로 작은 값부터 얻을 수 있도록 구성하고, 보석의 가격만 저장하는 배열을 따로 만들어 최대힙으로 구성하여 큰 값부터 얻을 수 있도록 하는 것이 핵심이었습니다.

이상입니다.