본문 바로가기

백준 그리디 알고리즘

파이썬 그리디 알고리즘 백준 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)

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

출력

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

풀이

보석 가격의 합이 가장 크도록 보석을 가방에 담아야 하는데

시간 초과가 나지 않게 하기 위해서 heapq를 이용해야 한다.

heapq는 리스트를 순회하지 않고 가장 작은 원소의 값을 구할 수 있다. (heapq[0][0], heapq.pop)

보석에 대한 무게, 가격 정보 리스트를 heapq로 만든다.

우선 보석의 정보를 보석의 무게로 오름차순 정렬하고 가방도 담을 수 있는 최대무게로 오름차순 정렬한다.

그리고 가방에 대해 반복문을 돌려서 보석 heapq리스트가 남아있고, 가장 작은 보석의 무게가 가방의 수용 무게

이하 일 경우 가방에 보석을 담을 수 있기 때문에 possible이라는 heapq리스트에 추가한다.

이때 보석은 무거운 순으로 담아야 하지만 heapq는 가장 작은 원소만 구할 수 있기 때문에 -기호를 붙혀서 추가한다.

한 가방에 대해 담을 수 있는 보석 무게를 possible 리스트에 전부 저장한 후에 heapq.pop을 통해

가장 무거운 보석 무게 1개만 money에 추가한다. (무게가 음수 기호기 때문에 -= 연산을 사용한다.)

모든 가방에 대한 반복을 수행한 후 최종 출력 값인 money를 출력한다.

 

import sys
import heapq

n, k = map(int, sys.stdin.readline().rstrip("\n").split())
nheaps = []
kLists = []
money=0
for i in range(n):
    m, v = map(int,sys.stdin.readline().rstrip("\n").split())
    heapq.heappush(nheaps,(m,v))

for i in range(k):
    kLists.append(int(sys.stdin.readline().rstrip("\n")))
kLists = sorted(kLists)

possible=[]
for i in range(0,len(kLists)):
    while nheaps and kLists[i]>=nheaps[0][0]:
        (m,v) = heapq.heappop(nheaps)
        heapq.heappush(possible,-v)
    if possible:
        money-=heapq.heappop(possible)
print(money)