백준 BOJ 1202 보석 도둑

문제

 

코드

import sys
import heapq

input = sys.stdin.readline
inf = sys.maxsize

if __name__ == "__main__":
    n, k = map(int, input().split())
    gems = [[*map(int, input().split())] for _ in range(n)]
    bags = [int(input()) for _ in range(k)]
    
    gems.sort() # 보석 무게 오름차순
    bags.sort() # 가방 용량 오름차순

    heap = []
    gem_index = 0
    result = 0
    
    for bag in bags:
        while gem_index < n and gems[gem_index][0] <= bag:
            heapq.heappush(heap, -gems[gem_index][1])
            gem_index += 1
        
        if heap:
            result += -heapq.heappop(heap)

    print(result)

 

문제 목표: 각 가방에는 보석 1개만 넣을 수 있다. 각 보석도 1번만 쓸 수 있다. 총 가격을 최대로 만들어라.

 

작은 가방부터 보면서, 그 가방에 들어갈 수 있는 보석들을 후보로 모으고, 그중 가장 비싼 보석을 선택한다.

gems는 [무게, 가격] 이고, bags 는 가방 용량이다.

gems.sort()
bags.sort()

정렬을 통해 작은 가방부터, 보석의 무게가 가벼운 것부터 순차적으로 배열이 가능하다.

작은 가방부터 보아서 넣을 수 있는 최대한의 보석을 가져가야 하기 때문이다.

 

이것은,

heap = []
gem_index = 0
result = 0

 

아래와 같다.

heap       = 현재 가방에 들어갈 수 있는 보석들의 가격 후보
gem_index  = 아직 확인하지 않은 보석의 위치
result     = 선택한 보석 가격의 합

 

heap에는 보석의 '가격'만 넣을 뿐이며 heapq는 최소힙이기 때문에 가장 큰 가격을 먼저 꺼내려면 - 를 넣어야 한다.

heapq.heappush(heap, -가격)

 

예:

heap = [-100, -70, -30]

 

이제 작은 용량의 가방부터 하나씩 처리한다.

for bag in bags:

 

아직 확인하지 않은 보석 중에서 현재 가방에 들어갈 수 있는 보석을 전부 heap에 넣는다.

while gem_index < n and gems[gem_index][0] <= bag:

 

예를 들어 현재 가방 용량이 2이고, 보석이 다음과 같다면

gems = [[1, 65], [2, 99], [5, 23]]

 

[1, 65] 가능
[2, 99] 가능
[5, 23] 불가능

 

=> 가격 65, 99를 heap에 넣는다.

 

heapq.heappush(heap, -gems[gem_index][1])
gem_index += 1

 

보석 가격을 힙에 넣고, 다음 보석으로 넘어간다.

 

그런 이후, 현재 가방에 넣을 수 있는 후보 중 가장 비싼 보석을 하나 꺼내고, 해당 보석은 사용한 것이므로 힙에서 제거한다.

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

 

예시

보석:
무게 1, 가격 65
무게 5, 가격 23
무게 2, 가격 99

가방:
2
10

 

정렬 후:

gems = [[1, 65], [2, 99], [5, 23]]
bags = [2, 10]

 

처음:

heap = []
gem_index = 0
result = 0

 

첫 번째 가방인 bag = 2

무게 1 보석 가능 → heap에 -65
무게 2 보석 가능 → heap에 -99
무게 5 보석 불가능 → stop

 

상태:

heap = [-99, -65]
gem_index = 2

 

그 중 가장 비싼거 선택:

heappop(heap) = -99
result += 99

 

상태:

result = 99
heap = [-65]

 

이제 두번째 가방인 bag = 10

gem_index = 2 부터 본다.

무게 5 보석 가능 → heap에 -23

 

상태:

heap = [-65, -23]

가장 비싼거 선택:

heappop(heap) = -65
result += 65

 

최종:

result = 164