1202. 보석 도둑 (Python)

1202번: 보석 도둑

 

풀이

 보석과 가방을 담아두는 리스트를 만들고 두 리스트를 오름차순으로 정렬한다(보석의 경우 무게를 기준으로). 문제에서 주어지는 N과 K의 최댓값이 300,000이기 때문에 단순한 방법으로 2중 for문을 사용할경우 시간초과가 발생한다. 

 

 시간초과를 방지하기 위해 가방 리스트를 기준으로 for문을 시작한다. 다음으로 보석 리스트에 대하여 for문을 돌면서 현재 가방에 담을 수 있는 무게 이하인 보석들을 우선순위 큐에 집어넣는다. 만약 초과하는 무게를 가진 보석을 만날 경우 break를 통해 for문을 빠져나온다. 그 후 우선순위 큐가 비어있지 않을 경우 가장 높은 가격을 가진 보석 하나를 꺼낸다. 즉, 현재 가방에 담을 수 있는 보석들 중 가장 가치가 높은 보석을 담는다는 의미이다. 

 

 이렇게 가방을 다 탐색하고 나면 훔칠 수 있는 보석 가격의 합의 최댓값을 구할 수 있다. 

 

코드

import sys
import heapq

input = sys.stdin.readline

if __name__ == '__main__':
    N, K = map(int, input().split())  # N: 총 보석 수, K: 가방 개수
    jewel = []  # 보석 리스트
    bags = []  # 가방 리스트

    for _ in range(N):
        M, V = map(int, input().split())  # M: 보석의 무게, V: 보석의 가격
        jewel.append((M, V))

    for _ in range(K):
        bags.append(int(input()))

    jewel.sort()  # 보석의 무게를 기준으로 오름차순 정렬
    bags.sort()  # 가방에 담을 수 있는 최대 무게 기준으로 오름차순 정렬
    res = 0  # 훔칠 수 있는 보석의 최대 가격
    pq = []  # 현재 가방에 담을 수 있는 보석들을 담아놓는 우선순위 큐

    idx = 0  # jewel 리스트의 인덱스
    for bag in bags:
        for i in range(idx, N):  # 처음부터가 아닌 idx부터 시작한다.
            if jewel[i][0] <= bag:  # 현재 가방에 담을 수 있는 보석이라면
                heapq.heappush(pq, -jewel[i][1])  # 우선순위 큐에 담는다.
                idx += 1 # 다음 가방에 담을 수 있는 보석을 찾을 때, jewel의 처음이 아닌 idx부터 찾을 수 있도록 idx 값 조정
            else:
                break

        if len(pq) != 0:  # 현재 가방에 담을 수 있는 보석이 있을 경우
            res += -heapq.heappop(pq)  # 가장 높은 가격을 가진 보석 하나를 가방에 담는다.

    print(res)