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

백준 1202 보석 도둑

숫자가 굉장히 크고 우선순위가 정해졌기 때문에, 그리디알고리즘을 활용해야 한다는 것은 바로 알 수 있었다. 내 오답은 솔루션 이후에 정리하겠다. 이 문제의 포인트는 다음과 같았다.

1. 가방과 보석을 오름차순 정렬한다
2. 작은 가방부터 지금까지 들어갈 수 있는 보석 중에 가장 가치가 높은 보석을 가방에 넣는다.
3. 가방을 모두 채우거나 보석이 없다면 반복문을 종료한다.

즉, 가방의 크기별로 정렬한 이후에 가방에 들어갈 수 있는 보석들을 저장할 리스트를 만든다. 이 리스트는 각 가방에 들어갈 수 있는 보석들이 아니라 지금까지 들어갈 수 있는 보석들이다. 이를 구현하는 과정에서 한 가지 실수를 했다. 보석들을 pop 하는 과정에서 heapq.heappop(jewelry) 대신 pop(0)을 사용했다가 시간 초과가 났다.

pop(0)는 O(1)이지만 리스트를 shift하는 과정 때문에 결과적으로는 O(n)이라는 결과를 낳지만, heapq.heappop(jewelry) 힙 구조를 사용하기 때문에 O(log n)의 시간 복잡도를 갖는다.

import sys
import heapq

n , k = map(int, sys.stdin.readline().split())

jewelry = [list(map(int, sys.stdin.readline().split())) for _ in range(n)]
bags = [int(sys.stdin.readline().strip()) for _ in range(k)]
tmp = []

jewelry.sort()
bags.sort()

res = 0
for bag in bags:
    while jewelry and bag >= jewelry[0][0]:
    	# pop(0)하면 시간 초과
        heapq.heappush(tmp, -heapq.heappop(jewelry)[1])
    if tmp:
        res -= heapq.heappop(tmp)
    elif not jewelry:
        break
print(res)

내 오답은 보석의 가치를 기준으로 내림차순 정렬하고 이후에 가방에 들어갈 수 있는 최적의 값을 찾는 것이었다. 최적의 가방을 찾는 과정에서 이분 탐색을 활용했지만 7%에서 시간 초과가 났다.

import sys
from queue import PriorityQueue

n , k = map(int, sys.stdin.readline().split())
cList = []
pq = PriorityQueue()
for _ in range(n):
    w, v = map(int, sys.stdin.readline().split())
    pq.put([-v, -w])

for _ in range(k):
    cList.append(int(sys.stdin.readline().strip()))

res = 0
cList.sort()
while cList and not pq.empty():
    v, w = [-x for x in pq.get()]
    # print(v, w)
    if w > cList[-1]:
        continue
    lp = 0
    rp =len(cList) - 1
    while lp < rp:
        mid = (lp + rp) // 2
        if cList[mid] < w:
            lp = mid + 1
        else:
            rp = mid
    cList.pop(lp)
    res += v
print(res)