백준 1202: 보석 도둑(골드 II) - Greedy, Priority Queue

1202번: 보석 도둑

유명한 보석 문제이다. 스터디에서 진행했던 문제이고 그리디 알고리즘을 공부할 수 있는 좋은 문제인 것 같아, 정리차원에서 포스팅한다.

문제 요약:

최대 용량이 다양한 가방들이 주어지고 가방에는 하나의 보석만 담을 수 있다.

보석은 무게와 값어치가 각각 주어지고, 최대 30만개를 훔칠 수 있다.

보석의 무게가 가방의 최대 용량보다 크면 가방에 담을 수 없다.

최대로 훔칠 수 있는 보석의 값어치 총합을 구하여라.

 

솔루션:

단순하게 생각하면 반례가 반드시 생기는 문제였던 것 같다. 해당 문제는 정렬, 그리디, 우선순위 큐 알고리즘이 모두 사용된다. 용량이 가장 큰 가방에 값어치가 가장 큰 걸 담는다면, 가방이 낭비되는 경우이다.

 

왜냐하면 무게가 30, 값어치가 1억인 보석과 무게가 10 값어치가 3억인 보석이 있고, 가방 용량이 100과 20인 두 개의 가방이 있다고 할 때, 100짜리 가방에 3억짜리 무게 10 보석을 넣는다면 값어치가 1억인 보석은 20용량 가방에 담을 수 없으므로 못 가져가게 된다. 반대의 경우는 모든 보석을 챙길 수 있다.

어떻게 하면 좋을까? 아래의 흐름대로 실행한다.

1. 가방을 무게 순으로 오름차순 정렬을 한다.

2. 보석을 무게 순으로 오름차순 정렬을 한다.

3. 가장 작은 가방에 담을 수 있는 모든 보석의 값어치를 리스트에 담는다.

4. 만약 보석의 무게가 가방 무게보다 크다면 그 뒤 보석들은 오름차순 정렬되어있기 때문에 볼 필요 없다.

5. 담긴 리스트 중 가장 값어치가 큰 보석을 더하고 리스트에서 제외한다.

6. 다음 가방으로 이동한다. 이때, 담긴 리스트는 그대로 유지된다.

 

자, 여기서 30만개의 보석이 주어지고, 가방이 바뀔 때 마다 최대30만개의 리스트에서 최댓값을 뽑으려면 O(N)이 소요되므로 O(n^2)이 소요된다. 반드시 시간초과가 발생한다. 

따라서 최댓값을 O(logN)에 뽑아낼 수 있는 우선순위 큐를 사용하며, 이를 위해 힙 자료구조를 사용한다. 즉,

3. 가장 작은 가방에 담을 수 있는 모든 보석의 값어치를 리스트 우선순위 큐에 담는다.

5. 담긴 리스트 중 가장 값어치가 큰 보석을 더하고 리스트 우선순위 큐에서 제외한다.

6. 다음 가방으로 이동한다. 이때, 담긴 리스트 우선순위 큐는 그대로 유지된다.

 

소스코드:

import sys
import heapq
from collections import deque
input = sys.stdin.readline

N, K = map(int,input().split())
gem_lst = []
back_lst = []
hq = []
result = 0
for _ in range(N):
    M, V = map(int, input().split())
    gem_lst.append((M, V))

back_lst = [int(input()) for _ in range(K)]
gem_lst.sort(key=lambda x: x[0]) # 무게순 정렬
back_lst.sort()

gem_lst = deque(gem_lst)

for back in back_lst:
    while gem_lst:
        m, v = gem_lst[0]
        if m <= back:
            heapq.heappush(hq, -v)
            gem_lst.popleft() 
        else:
            break 

    if hq:
        result += -heapq.heappop(hq)

print(result)