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

https://www.acmicpc.net/problem/1202

1202번: 보석 도둑

---

아이디어

1. 보석을 가방에 어떻게 담을 건가.

• 단순히 이중 반복문으로 비교하면 시간 초과가 난다.
• 다시 문제를 읽고 문제에 있는 힌트를 보고 감을 잡았다.
• 먼저 보석 리스트와 가방 리스트를 오름차순으로 정렬하고
• 각 가방을 기준으로 보석이 가방에 넣을 수 있는지 확인한다.
• 가방에 넣을 수 있는 보석 중 가치가 가장 큰 걸 넣어주면 된다.

2. 시간 단축

• 근데 그냥 반복문으로 작성하면 시간 초과가 또 발생한다.
• 가장 가치가 큰 것만 뽑아주면 되니까 최대 힙을 사용하기로 했다.

 

---

전체 코드

import sys
import heapq

N, K = map(int,input().split())
jewel_list = []
pocket_list = []

for _ in range(N):
    jewel_list.append([int(x) for x in sys.stdin.readline().rstrip().split()])
for _ in range(K):
    pocket_list.append(int(sys.stdin.readline().rstrip()))
jewel_list.sort()
pocket_list.sort()

ans = 0
can_put_in = []

for pocket in pocket_list:
    while jewel_list:
        if pocket >= jewel_list[0][0]:
            heapq.heappush(can_put_in, -jewel_list[0][1])
            heapq.heappop(jewel_list)
        else:
            break

    if can_put_in:
        ans += -heapq.heappop(can_put_in)
    else:
        if not jewel_list:
            break

print(ans)

 

---

코드 설명

for _ in range(N):
    jewel_list.append([int(x) for x in sys.stdin.readline().rstrip().split()])
for _ in range(K):
    pocket_list.append(int(sys.stdin.readline().rstrip()))
jewel_list.sort()
pocket_list.sort()

먼저 보석과 가방의 정보를 입력받고 둘 다 오름차순으로 정렬해준다.

 

ans = 0
can_put_in = []

for pocket in pocket_list:
    while jewel_list:
        if pocket >= jewel_list[0][0]:
            heapq.heappush(can_put_in, -jewel_list[0][1])
            heapq.heappop(jewel_list)
        else:
            break

    if can_put_in:
        ans += -heapq.heappop(can_put_in)
    else:
        if not jewel_list:
            break

print(ans)

 

가방이 작은 순서대로(이미 오름차순으로 정렬) 보석이 들어가는지 확인해본다.

가방 허용 무게보다 작거나 같은 보석들은 can_put_in리스트에 넣어준다.

그리고 jewel_list에서 보석을 빼준다. 

도중에 보석이 가방에 넣는 게 불가능하면 break 해준다.

 

이후 can_put_in 리스트가 비어있지 않으면 pop 해준다. 

can_put_in리스트는 최대 힙과 같은 형태이므로 처음에 있는 값이 가치가 가장 큰 값이다.

 

can_put_in 리스트도 비어있고 jewel_list도 비어있으면 break 해준다.(더 이상 넣을 게 없다는 뜻)