[Python] 백준 1202 보석 도둑

🧑‍💻 백준 1202 보석 도둑

Gold 2 - 정렬, heapq

 

각 보석들의 무게와 가치가 주어지면서, 각 가방에 들어갈 수 있는 보석의 최대 무게가 주어진다

각 가방에는 보석 하나 밖에 못 들어간다

 

즉, 가방에 들어갈 수 있는 보석들 중, 최대의 보석 가치를 가방에 넣는 것이다

---

 

문제 풀이

가방에 들어갈 수 있는 보석들 중, 최대의 보석 가치를 가방에 넣는 것

• 즉 보석들을 각 보석들의 무게 위주로 오름차순으로 정렬을 한다
• 그리고 가방에 들어갈 수 있는 보석들을 구한다
• 그 가방에 들어갈 수 있는 보석들 중, 제일 가치가 높은 보석을 가방에 넣으면 된다

heapq

• 힙은 이진 트리를 사용하면서 우선순위 큐를 구해준다.
• O(logN)

---

 

코드

결과 코드

import heapq

N, K = map(int, input().split())

treasure = []
bag = []

for _ in range(N):
    treasure.append(tuple(map(int, input().split())))

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

treasure.sort()
bag.sort()

temp = []
result = 0

for b in bag:

    while treasure and treasure[0][0] <= b:
        heapq.heappush(temp, -heapq.heappop(treasure)[1])

    if temp:
        result += heapq.heappop(temp)

print(-result)

 

 

코드 별 정리

temp = []

• 가방에 들어갈 수 있는 보석의 가치를 넣는 리스트다
• bag 리스트는 오름차순으로 정렬이 되어 있다
  • 즉 temp 리스트 안에 존재하는 보석의 가치들은, 모든 가방에 넣을 수 있는 보석의 가치들이다
  • 예) bag = [2, 10] / treasure = [(2, 55), (2, 65), (3, 99), (5, 23)]
    • bag에서 2를 순회했을 때, temp에는 [65, 55]이 들어가 있음. treasure의 첫 번째와 두 번째 보석의 가치
      • heap을 사용했기 때문에, 큰 숫자부터 정렬이 된다
      • 즉 65가 여기서 빠져 나간다.
    • bag에서 10을 순회를 했을 때, temp에는 [99, 55, 23] 이 있다
      • 남아있는 보석들의 무게는 모두 10을 안 넘어서, temp 안에 다 들어가게 된다.
      • temp를 통해, 제일 큰 가치의 보석을 가지고 가면 된다

heapq.heappush(temp, -heapq.heappop(treasure)[1])

• heapq.heappush(temp, -heapq.heappop(treasure)[1])
  • temp 리스트 안에, -heapq.heappop(treasure)[1]라는 원소를 넣는 것이다

  • -heapq.heappop(treasure)[1]

    • treasure 리스트의 첫 번째 원소의 가치를 음수로 바꾼다 (최대 힙으로 만들기 위해서)
    • 만약 [5, 3, 2, 6]이 있다
      • 그냥 heapq을 써서 한 숫자씩 꺼냈을 경우
        • 2, 3, 5, 6 순으로 출력이 된다
      • 반대로 음수로 출력했을 때에는
        • -6, -5, -3, -2 순으로 출력이 된다
        • 위에 숫자들을 다시 양수로 바꿔주면 내림차순으로 출력이 된다
• heapq.heappush()
  • temp 리스트 안에 넣는데, heappush를 썼으니, temp에서 제일 앞에 있는 숫자는, 리스트 중 제일 큰 숫자가 된다

result += heapq.heappop(temp)

• 즉 temp 안에 숫자가 있으면, 숫자들 중, 제일 앞에 있는 숫자를 result에 누적을 시켜주면, 답이 나온다

 

실패 코드

from collections import deque
import sys

N, K = map(int, sys.stdin.readline().split())

treasure = []
bag = []
result = 0

for _ in range(N):
    treasure.append(tuple(map(int, sys.stdin.readline().split())))

for _ in range(K):
    bag.append(int(sys.stdin.readline()))


treasure.sort(key=lambda x : -x[1])
bag.sort(key=lambda x: x)

for b in bag:
    i = 0

    while i < len(treasure):
        if treasure[i][0] <= b:
            result += treasure[i][1]
            treasure.pop(i)
            break

        elif treasure[i][0] > b:
            i += 1

print(result)

• while문을 두 번이나 겹치게 쓰면서, elif문으로 가면, 지속적으로 탐색하게 만들어서 시간 초과가 난 것 같다...