[boj] 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 1
5 10
100 100
11

예제 출력 1

10

예제 입력 2

3 2
1 65
5 23
2 99
10
2

예제 출력 2

164

풀이

처음에는 간단한 정렬 문제인 줄 알았는데 시간 초과로 상당히 고생했다.

처음 접근법 - 시간 초과

처음에는 높은 가치를 가진 보석을 기준으로 가장 가벼운 무게의 가방을 찾아서 넣으려고 했다. 이렇게 되면 각 보석마다 넣을 수 있는 가방까지 순회를 하면서 찾아야 하므로 시간 복잡도가 O(N^2) (N과 K의 최댓값이 동일함)이 되어 시간 초과로 실패를 했다.

import sys
import heapq

input = sys.stdin.readline
N, K = list(map(int, input().split()))
# 보석 (무게, 가치) 리스트 - 가치 순으로 정렬
gems = sorted([tuple(map(int, input().split())) for _ in range(N)], key=lambda x: x[1])
# 가방 무게 제한 리스트
bags = sorted([int(input()) for _ in range(K)])

answer = 0
candidateBags = []

while bags and gems:
    weight, value = heapq.heappop(gems)
    for i in range(len(bags)):
        if bags[i] >= weight:
            bags = bags[:i] + bags[i + 1 :]
            answer += value
            break
print(answer)

성공한 풀이

성공한 풀이에서는 현재 가방에 넣을 수 있는 보석이면 이 가방보다 더 많은 무게를 견딜 수 있는 가방에도 넣을 수 있다는 사실을 파악하는 것이 가장 중요했다. 이것을 확인한 후로는 하위 문제들의 최적의 해를 바탕으로 최종적인 답을 도출해낼 수 있는 구조이므로 그리디를 통해서 문제를 풀 수 있다.

이를 통해서 가장 작은 가방부터 확인을 하면 한 번 확인한 보석을 다시 가방에 넣을 수 있는지 없는지 확인하지 않아도 되므로 loop 안에 시간 복잡도가 O(N + K)가 되어 정렬에 드는 시간이 가장 큰 영향을 미치게 되어 최종적인 시간 복잡도가 O(NlogN) 되어 시간 초과 없이 문제를 해결할 수 있게 된다.

import sys
import heapq

input = sys.stdin.readline
N, K = list(map(int, input().split()))
# 보석 (무게, 가치) 리스트 - 무게 순으로 정렬
gems = sorted(
    [tuple(map(int, input().split())) for _ in range(N)],
    key=lambda x: x[0],
    reverse=True,
)
# 가방 무게 제한 리스트
bags = sorted([int(input()) for _ in range(K)])

answer = 0
candidateGemValues = []
for bag in bags:
    while gems and bag >= gems[-1][0]:
        _, value = gems.pop()
        # 현재 가방에 넣을 수 있는 높은 가치 순으로 정렬
        heapq.heappush(candidateGemValues, -value)
    if candidateGemValues:
        # 이전 가방은 현재 가방보다 제한이 낮았으므로
        # 이전 가방에 넣을 수 있는 보석은 현재 가방에도 넣을 수 있음
        answer += -1 * heapq.heappop(candidateGemValues)
print(answer)