[코딩 테스트 연습(파이썬/Python)] 백준 1202번 _ 보석 도둑

https://www.acmicpc.net/problem/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) 넣을 수 없는 경우에는 pass, 가방을 모두 사용하거나 모든 보석에 대해 시도하면 종료한다.

 

# 시간 초과
# import sys
# input = sys.stdin.readline
n,k = map(int,input().split())
j = [list(map(int,input().split())) for _ in range(n)]
b = [int(input()) for _ in range(k)]
j.sort(key=lambda x: x[1])
b.sort()
def put(j,b):
    bb=len(b)
    def find(m,b,bb):
        i=0;e=bb-1
        if m> b[e]:
            return 'No'
        while i!=e:
            mid = (i+e)//2
            if m==b[mid]:
                return mid
            elif m>b[mid]:
                i=mid+1
            else:
                e=mid
        return i
    answer = 0
    while bb and j:
        m,v = j.pop()
        k = find(m,b,bb)
        if k!='No':
            b.pop(k)
            answer +=v
            bb-=1
    return answer            
        
print(put(j,b))

시간 초과가 발생하는 이유로 추정되는 것은 아래와 같다.

 

① 모든 보석에 대해 이진탐색을 해야함

② 사용한 가방을 목록에서 제거하는 b.pop() 

 

해당 문제들을 해결하기 위해서는, 다른 접근이 필요하다.

 

[아이디어(두번째 시도 - 성공)]

1) 가방은 작은 순서대로, 보석은 가벼운 순서대로 정렬한다.

2) 가방을 하나씩 꺼내면서 해당 가방에 담을 수 있는 모든 보석을 순서대로 담은 목록을 형성한다.

3) 목록이 형성되면, 해당 가방에서 가장 가치 있는 보석은 실제로 담은 후, 나머지 목록(heapq)은 다음 가방에 넣을 수 있도록 유지한 채로 다음 가방으로 넘어간다.

4) 다시 새로운 가방에 담길 수 있는 모든 보석을 목록에 추가한 뒤, 가장 가치 있는 것을 실제로 담고 다음 가방으로 넘어간는 것을 반복한다.

5) 가방을 다 사용하거나, 보석을 모두 탐색했다면 종료한다.

 

import sys
input = sys.stdin.readline
import heapq
n,k = map(int,input().split())
j = [list(map(int,input().split())) for _ in range(n)]
tmp = []
b = [int(input()) for _ in range(k)]
j.sort(key=lambda x: x[0])
b.sort()
answer = 0
i = 0
for bag in b:
    while i < n and bag >= j[i][0]:
        heapq.heappush(tmp, -j[i][1])
        i += 1
    if tmp:
        answer += heapq.heappop(tmp)
print(-answer)

벽을 느끼고 다른 사람의 코드를 참고했다^ 일종의 투포인터를 사용한 것이나 다름 없다.(포인터1: 가방, 포인터2: 보석) 매번 이진 탐색을 진행하지 않기 위해 순서가 유지되는 heapq를 사용하고, 최대힙화하기 위해 음수를 사용했다. 가장 가치가 높은 보석부터 담기 시작해야한다고 생각했던 것이 문제였다. 각 가치의 보석을 담을 최소의 가방에 담기게 하면 되는 것이지, 가치가 높은 것부터 시도할 이유는 따로 없다.