문제 링크 : https://www.acmicpc.net/problem/1202
생각 바운더리를 한 단계 넓혀준 중요한 문제인 것 같다. 문제는 단순한데 풀이 방법은 단순하지 않다. 내 기준 그리디 알고리즘 대표 문제로 뽑아도 될 만큼 좋은 문제라고 생각했다.
그리디 문제를 풀 때, 정렬을 함으로써 생기는 추가적인 조건이 생긴다는 것을 잊지 말아야겠다. 이 문제에선 가방을 오름차순으로 정렬했을 때, 현재 가방의 다음 순서에 오는 가방들은 항상 현재 가방이 수용할 수 있는 보석들을 담을 수 있다는 거다. 생각해보면 당연한 거지만 이걸 코드에 적용을 해서 풀어야 시간복잡도 효율을 높일 수 있다.
처음에 이 문제를 풀 때 그리디적으로 "최대한 작은 가방으로 최대한 비싼 보석을 담는다" 는 생각은 할 수 있었다.
그래서 보석은 비싼 순으로, 가방은 가벼운 순으로 정렬했다. 그리고 보석의 입장에서 생각을 했다. 예를 들어 비싼 순으로 보석을 정렬했을 때, 인덱스 0 보석의 무게가 10 가격이 100 이라고 했을 때, 보석에서 for 문을 돌릴 생각을 했다.
현재 제일 비싼 보석이 무게가 10이기 때문에, 나는 이 보석을 담기 위해 가방중에서 무게가 10이상인 가장 가벼운 가방을 탐색해야 한다. 이 과정이 최악 O(K) time 이 소요되는데, N 과 K 의 범위는 300,000 이기 때문에 O(N * K) time 은 시간 초과로 문제를 틀리게 된다.
그래서 위에서 말했던 조건을 꼭 적용해줘야된다. 그리고 보석의 입장보다는 가방의 입장에서 생각하는게 좀 더 자연스럽다. 왜냐하면 결국 문제가 요구하는 건 도둑이 취할 수 있는 최대 가격의 보석이기 때문이다. 보석은 안 담는 보석도 많다. 그치만 가방은 가급적 다 사용해야 한다.
가방도 가벼운 순으로, 보석도 가벼운 순으로 정렬해야 된다. 그래서 가방을 기준으로 for 문을 돌려야 된다. 현재 가방에 담을 수 있는 무게의 보석들을 while 문을 이용해서 전부 담는다. 담을 때는 가격을 기준으로 한 Max Heap (candidate) 에 담는다. 전부 담은 뒤에, 그 중에서 제일 비싼 보석만 뽑아낼 거기 때문이다 (heappop).
현재 가방 처리를 마친 뒤에 다음 가방을 처리할 때는 "조건"에 의해서 다시 처음부터 보석을 탐색할 필요가 없다. 이전 가방에서 탐색했던 보석들은 전부 현재 가방에 담을 수 있는 보석들이다. 전에 탐색했던 마지막 보석의 다음 위치 보석부터 탐색을 진행 하면 된다.
그래서 이렇게 하면 위에서 했던 풀이와 다르게 O(N + K) time 이 소요되기 때문에 시간 초과없이 정답을 맞출 수 있게 된다.
파이썬 heapq 클래스를 import 해야되고, 파이썬의 힙은 기본이 minHeap 이기 때문에 (-가격) 을 heappush 해줘야 maxHeap 처럼 사용할 수 있다.
import sys import heapq N, K = map(int, sys.stdin.readline().split()) gems = [] bags = [] for _ in range(N): M, V = map(int, sys.stdin.readline().split()) gems.append((M, V)) for _ in range(K): C = int(sys.stdin.readline()) bags.append(C) # 가벼운 보석 순으로 정렬 gems.sort() # 가벼운 가방 순으로 정렬 bags.sort() candidate = [] i = 0 answer = 0 for bag in bags: while i < len(gems) and bag >= gems[i][0]: heapq.heappush(candidate, -gems[i][1]) i += 1 if candidate: maxGem = heapq.heappop(candidate) answer += maxGem * (-1) print(answer)