백준 1202번: 보석 도둑 [python]

난이도 : 골드 2
유형 : 우선순위 큐
https://www.acmicpc.net/problem/1202

문제 접근

보석 N개: 각 보석은 무게 M, 가격 V를 가짐.
가방 K개: 각 가방은 최대 무게 C까지 수용 가능.
가방 하나에는 보석 하나만 넣을 수 있음.
훔칠 수 있는 보석들의 최대 가격 합을 구하라. 라는 문제.

해결 아이디어

해당 가방에 들어갈 수 있는 모든 보석을 우선순위 큐에 넣음 (가격 기준 max-heap)
그 중 가장 가격이 높은 보석(max-heap[0])을 꺼내어 사용하는 논리로 풀었다.

해답 및 풀이

import heapq
import sys
input = sys.stdin.readline

# 1. 입력

# N : 보석 개수, K : 가방 개수
N, K = map(int,input().split())


# 보석 정보 입력 튜플로 저장, 무게와 가치를 따로 구분하기 위해
jewels = []
for _ in range(N):
    m, v = map(int,input().split())
    jewels.append((m, v))

# 가방
bags = []
for _ in range(K):
    c = int(input())
    bags.append(c)

jewels.sort() # 무게 기준 오름차순 정렬
bags.sort()

result = 0
jewel_idx = 0
heap = [] # max-heap 으로 사용 (가격 기준))

for bag_weight in bags:
    # 현재 가방에 넣을 수 있는 모든 보석들을 heap에 추가.
    while jewel_idx < N and jewels[jewel_idx][0] <= bag_weight:
        
        heapq.heappush(heap, -jewels[jewel_idx][1])
        jewel_idx += 1

        # 가장 비싼 보석 꺼내 가방에 넣기

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

코드를 보고 이해는 되는데 매번 내 손으로 써내려갈때 입력문 다음에서 바로 막힌다.