[그리디 알고리즘4] 백준 1202 보석 도둑 파이썬 우선순위 큐 heapq

백준 1439 보석 도둑 

 

📜 문제 

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

상덕이가 털 보석점에는 보석이 총 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)

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

 

 

🌞 출력

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

 

 

🍖문제 풀이 

 

 

🥕 정답 코드 

출력 초과가 나온 코드 

출력 초과란 ? = 틀렸다! 정답 파일 크기의 2배를 넘어가면 '출력 초과' 를 받게 된다고 한다. ex. 무한 루프 등.. 

다시보니 print로 디버깅한 곳들 때문이었다. 필요없는 print를 지우니 ' 시간 초과' 가 나왔다. -> import sys해서 다시 제출하니 '틀렸습니다' 나왔다.. ㅎ..결국 틀린 것이었.......

n, k = map(int,input().split())

gem_info = []
for i in range(n):
    x,y = map(int, input().split())
    gem_info.append([x,y])
print(gem_info)

gem_info.sort(key=lambda x: x[0])
gem_info.sort(key=lambda x: x[1], reverse=True)
print(gem_info)

k_w = []
for i in range(k):
    x = int(input())
    k_w.append(x)
k_w.sort(reverse=True)

price = 0 
index = k
answer = 0 
for i in range(k):
    for j in range(len(gem_info)):
        if gem_info[j][0]<=k_w[i]:
            price += gem_info[j][1]
            
            index -= 1
            if index==0: break
    break

print(price)

 

찾아보니 이 문제는 '우선순위 큐' heapq을 사용해서 푸는 문제였다. heapq가 뭔지 오늘 알았다.... 유튜브로 공부를 하고 정답 코드를 보면서 이해했다. 

 

정답 코드

import sys
import heapq
input = sys.stdin.readline

n, k = map(int,input().split())

gem_list = [list(map(int,input().split())) for _ in range(n)]
bag_list = [int(input()) for _ in range(k)]
gem_list.sort()
bag_list.sort()

result = 0
temp = []

 #bag_list 가방이 담을 수 있는 무게 = i 
for i in bag_list: 
    # gem_list가 존재하고 가방이 담을 수 있는 무게가 보석의 무게와 같거나 클 때 
    while gem_list and i>=gem_list[0][0]:
        # temp에 보석 가격 입력 
        #- 를 붙여서 max heap 구현 
        heapq.heappush(temp, -gem_list[0][1])
        heapq.heappop(gem_list)

    # temp의 요소가 있으면 
    if temp:
        result += heapq.heappop(temp)
  


print(-result)

heapq의 시간복잡도는 O(logN) 이다 . temp 에다가 haaqp.push 를 사용해서 gem_list의 가장 첫번째(sort 하였으므로) 무게가 적은 보석의 가격 gem_list[0][1] 을 넣어준다. 그런데 python heapq는 min_heap이다 즉, pop을 하게 되면 가장 작은 값을 우선 주는 아이! 여기서 가격 최대값을 구해야하기 때문에 - 마이너스를 붙여서 넣어준다. 그러면 나중에 pop하게 될 때 가장 작은 숫자(마이너스를 붙인) 가 다시 - 마이너스를 붙여 + 로 만들면  최대 가격을 구할 수 있다. 

답을 print할 때 -result 부분이 요 부분. 

 

아래의 코드에서는 아예 result (or answer)에 최대 보석의 가격을 -= 로 더해버렸다 

 

똑같지만 더 짧은 코드

import sys
import heapq
input = sys.stdin.readline

n,k = map(int,input().split())
gem = [list(map(int,input().split())) for _ in range(n)]
bag = [int(input()) for _ in range(k)]
gem.sort()
bag.sort()

Q =[]
answer = 0
for i in bag:
    while gem and gem[0][0]<=i:
        heapq.heappush(Q, -heapq.heappop(gem)[1])
    if Q: answer -= heapq.heappop(Q)
print(answer)