[백준] 1202 - 보석 도둑 (Python)

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

1) 아이디어

• 보석과 가방을 무게순으로 내림차순 정렬한 후, 현재 가방에 들어갈 수 있는 보석 중에서 가장 가치가 높은 보석을 가치순으로 재정렬한다.
• 보석과 가방의 수가 항상 일치하지 않기 때문에, 아무것도 담기지 않은 가방이 발생하거나 보석이 남을 수 있으므로 이를 고려해 작성해야 한다.

2) 풀이

heap을 사용해 보석 정렬하기

a = []
for c in bags:
    while gems and ( gems[-1][0] <= c ):
        m, v = gems.pop()
        heappush(a, (-v, m))
    
    if a:
        v, m = heappop(a)
        ans += -v

입력받은 N개의 보석과 K개의 가방을 정렬한 후, 현재 가장 용량이 작은 가방부터 담을 수 있는 보석들을 힙 a에 삽입한다. a가 비어 있으면 현재 가방에 넣을 수 있는 보석이 존재하지 않으므로 가방을 건너뛰고, 그렇지 않은 경우에는 현재 넣을 수 있는 보석 중 가장 가치가 높은 보석을 인출해 정답에 더한다.

가방의 무게를 오름차순 정렬했기 때문에 다음 가방으로 넘어가더라도 현재 힙 a에 삽입된 보석들은 모두 다음 가방에 수납이 가능하므로 루프마다 a를 다시 선언하거나 재정렬할 필요가 없다.

참고) heapq

python의 heapq에서 제공하는 힙은 최소 힙으로, 부모 노드의 값이 자식 노드의 값보다 크도록 구성된다. 힙은 부모 노드와 자식 노드의 대소관계만 지킬 뿐, 오른쪽 자식노드와 왼쪽 자식노드간의 대소관계는 보장되지 않기 때문에 힙을 리스트처럼 순회하는 경우에는 정렬을 보장하지 않으므로 유의해야 한다.

다음은 10개의 랜덤한 정수에 대해 힙 h를 이용해 정렬된 리스트를 얻는 간단한 예시이다.

from heapq import heappop, heappush
import random

h = []
for _ in range(10):
	x = random.randint(1, 100)
    heappush(h, x)

sorted_list = []
while h:
	sorted_list.append(heappop(h))

힙에 대한 삽입/인출 시간복잡도는 $O(\log n)$이며, 기존 선언된 리스트 등을 힙으로 변환하는 heapify 연산은 $O(n)$이다. 특수한 경우가 아닌 경우 일반 정렬이 빠르지만, 이 문제와 같이 원소가 계속 삽입/인출되는 조건에서 정렬 상태를 유지해야 하는 경우 효율적이다.

3) 전체 코드

import sys
from heapq import heappop, heappush
input = sys.stdin.readline

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

gems = []
for _ in range(N):
    gems.append(tuple(map(int, input().split())))
gems.sort(reverse=True)

bags = []
for _ in range(K):
    bags.append(int(input()))
bags.sort()

ans = 0
a = []
for c in bags:
    while gems and ( gems[-1][0] <= c ):
        m, v = gems.pop()
        heappush(a, (-v, m))
    
    if a:
        v, m = heappop(a)
        ans += -v

print(ans)