[백준] 보석 도둑 (python)

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

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

 

아이디어

• 가방에 담을 수 있는 정보 저장 : 최소힙 이용
• 그 중에서 가장 비싼 보석 선택 : 최대힙 이용

 

풀이

 

<3번 시행착오를 거친 풀이 과정>

먼저, 가방과 보석의 무게, 가치 등을 정렬한 뒤에 가방 반복문을 돌면서 보석을 선택해야 된다고 생각했다.

여기서 생각해봐야 하는 점은 무엇을, 어떤 기준으로 정렬할 것이며, 보석은 어떻게 선택할 것인가

 

풀이1) 이중반복문, 시간초과

가장 쉽게 떠올릴 수 있는 방법은 2중 반복문이다.

1. 정렬
   - 가방은 오른차순 정렬
   - 보석은 무게 순으로 오른차순 정렬
2. 무게가 작은 가방부터 보석을 선택
   - 자신이 수용할 수 있는 무게의 보석까지 탐색
   - 그 중 가치가 최대인 것을 담는다.

이 방법으로 하면 O(n*k) ⇒ (1 ≤ N, K ≤ 300,000) 이므로

N*K는 900억,, 당연히 시간 초과이다.^^

 

풀이2) heapq, 시간초과

이런 그리디 문제에서는 주로 heapq를 사용해 시간을 줄였던 게 생각났다.

그래서 heapq를 어떻게 이용할지 고민해보았는데,

내가 생각한 두번째 풀이는 다음과 같다.

1. 정렬
   - 가방 오름차순 정렬
   - 보석은 (무게, 가치) 로 heapq 에 삽입한다. (즉, 무게 기준으로 오름차순 정렬)
2. 무게가 작은 가방부터 보석을 선택
   2-1) 가방에 넣을 수 있는 보석을 모두 pop한다. (최소힙 이용)
   2-2) possible이라는 heaqp에 (-가치)로 push한다. (최대힙 이용)
3. possible에 값이 있다면 그 중 가치가 최댓인 것을 pop하여 result에 더한다.
4. 남은 possible에 있는 보석들은 다시 jewels에 추가한다.

⇒ 하지만 이것도 시간초과가 난다.

시간복잡도를 어림잡아 계산해보면 최악의 경우 마찬가지로 O(N * K)이다.

heapq 삽입, 추가가 O(1)이라고 하더라고

2)번과 3)번에서 가방에 넣을 수 있는 보석의 개수가 K개 인 경우가 생길 수 있기 때문이다.

 

그래서 이 과정을 생략하고 풀 수 있는 방법을 고민해보았다.

 

 

풀이3) heapq, 성공

1. 정렬
   - 가방 오름차순 정렬
   - 보석은 (무게, 가치) 로 heapq 에 삽입한다. (즉, 무게 기준으로 오름차순 정렬)
2. 무게가 작은 가방부터 보석을 선택한다.
   2-1) while 조건식
       - 보석 heapq에 값이 있나?
       - 가장 앞 원소(즉, 무게가 가장 작은 보석)이 가방에 들어가는가?
   2-2) True라면 possible이라는 heapq에 (-가치) 로 push한다. (최대힙 이용)
   2-3) False라면 (즉, 더 이상 꺼낼 수 있는 보석이 없거나 / 가방에 들어가는 보석이 없거나) break
   ⇒ 결과적으로 해당 가방에 들어갈 수 있는 보석들은 모두 possible에 있다.
3. possible 에 보석 정보가 있다면
   3-1 ) 그 중에서 최대값을 pop한다.
   3-2 ) result 에 합한다.

여기서 possible에 있는 보석정보를 jewels에 다시 옮기지 않아도 되는 이유는

가방이 작은 순서대로 진행되고 있기 때문에

possible에 있는 보석정보는 어차피 나중에 진행되는 가방에도 넣을 수 있다!

 

+) 보석의 개수가 가방의 개수보다 적을수도 있다.

이 경우를 위해 possibles에도. jewels에 보석 정보가 없으면 break로 빠져나오는 코드를 추가한다.

코드

import sys, heapq
input = sys.stdin.readline
n, k = map(int, input().split())
jewels = []
for i in range(n):
    m, v = map(int, input().split())
    heapq.heappush(jewels, (m, v))
bags = []
for i in range(k):
    bags.append(int(input()))

bags.sort()
result = 0
possible = []
for bag in bags:
    while jewels and jewels[0][0] <= bag:
        heapq.heappush(possible, -heapq.heappop(jewels)[1])

    if possible:
        max_jewel = heapq.heappop(possible)
        result += -max_jewel
    elif not jewels:
        break


print(result)

 

 

알고리즘 유형

- 그리디