[백준] 1202(보석 도둑) 파이썬

문제

세계적인 도둑 상덕이는 보석점을 털기로 결심했다.

상덕이가 털 보석점에는 보석이 총 N개 있다. 각 보석은 무게 Mi와 가격 Vi를 가지고 있다. 상덕이는 가방을 K개 가지고 있고, 각 가방에 담을 수 있는 최대 무게는 Ci이다. 가방에는 최대 한 개의 보석만 넣을 수 있다.

상덕이가 훔칠 수 있는 보석의 최대 가격을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000)

다음 N개 줄에는 각 보석의 정보 Mi와 Vi가 주어진다. (0 ≤ Mi, Vi ≤ 1,000,000)

다음 K개 줄에는 가방에 담을 수 있는 최대 무게 Ci가 주어진다. (1 ≤ Ci ≤ 100,000,000)

모든 숫자는 양의 정수이다.

출력

첫째 줄에 상덕이가 훔칠 수 있는 보석 가격의 합의 최댓값을 출력한다.

풀이 및 리뷰

1. 보석과 가방을 크기 순으로 정렬한다.

2. 가장 작은 가방부터 넣을 수 있는 보석을 확인한다.

3. 넣을 수 있는 보석들을 모두 우선순위 큐에 넣는다.

4. 넣을 수 있는 보석이 있다면 가치가 가장 높은 것(우선순위 큐 pop)을 가방에 넣는다.

 

핵심은 우선순위 큐를 사용해 한번 확인한 보석은 다시 확인하지 않는 것입니다. 가방의 크기가 커지면 추가로 들어갈 수 있는 보석을 갱신하고, 그중에서 가장 가치가 큰 보석을 찾습니다.

 

우선순위 큐는 heapq를 통해 구현할 수있습니다.

import heapq

q = [] # 우선순위 큐로 사용할 리스트
q.heappush(q, 1) # 저장할 리스트와 저장할 요소를 인자로 넣음
q.heappush(q, 3)
q.heappush(q, 2)

q.heappop(q)
q.heappop(q)
q.heappop(q)
# 1, 2, 3

heapq는 요소를 넣을 때마다 리스트를 오름차순으로 정렬합니다. ( O(log(n) )

내림차순으로 정렬하고 싶을 경우엔 - 를 붙여 push하고 - 를 붙여 pop 하면 됩니다.

 

그리고 컴파일할 때 입력량이 많거나 계산량이 많다면 pypy3로 컴파일하면 python3에 비해 컴파일 시간이 줄어듭니다.

코드

import heapq
import sys
readline = sys.stdin.readline

N, K = map(int, readline().split())
gems = []
bags = []
for _ in range(N):
    m, v = map(int, readline().split())
    gems.append((m, v))
for _ in range(K):
    bags.append(int(readline()))

gems.sort()
bags.sort()
q = []
result = 0

i = 0
for bag in bags:

    while i < N and gems[i][0] <= bag:
        heapq.heappush(q, -gems[i][1])
        i += 1

    if q:
        result += -heapq.heappop(q)

print(result)

 

1202번: 보석 도둑