[백준 1202번] 보석 도둑 (Python3)

import sys
input = sys.stdin.readline
from heapq import heappush, heappop
    
N,K = map(int,input().split())

jewel = []
for i in range(N):
  heappush(jewel,tuple(map(int,input().split())))

bag = [0]*K
for i in range(K):
  bag[i] = int(input())
bag.sort()

result = 0
jewel1 = []
for weight in bag:
  while jewel:
    if jewel[0][0]>weight:
      break
    w,v = heappop(jewel)
    heappush(jewel1,(-v,w))
  if jewel1:
    v,w = heappop(jewel1)
    result += abs(v)
      
print(result)

 

어려웠던 문제다. 처음엔  단순히 보석을 value 순으로 나열하고, 이중 for문을 돌려서 value가 높은 순서대로 넣을 수 있는 가장 작은 가방에 넣으면 된다고 생각했다.

그러나 이 방식으로는 아무리 가지치기를 해도 시간초과를 해결할 수 없었다.

 

문제를 해결하기 위해서는 두 가지의 정렬기준이 필요하다.

1. 처음에는 보석을 무게 순으로 정렬하는 heapq에 담는다.

2. 가장 작은 배낭부터, 배낭의 용량보다 작은 보석을 두번째 heapq에 value 순으로 담는다. 이때 가치는 최대힙으로 담아야한다. 

3. 두번째 힙큐에서 가장 value가 높은 보석을 가방에 담는다.

 

이런 약간 복잡한 과정을 거쳐야 시간초과를 당하지 않는다.

두가지 heapq를 서로 다른 기준에 따라 이용한다는 것이 참신했던 문제였다.

 

 

오늘의 교훈) heapq를 잘 활용하자