Baekjoon 1202. 보석 도둑 / Python

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

문제요약

상덕이는 보석을 훔치려 한다. 각 보석은 무게와 가격이 있으며, 상덕이는 각 가방에 최대 한 개의 보석만 넣을 수 있다. 주어진 가방의 최대 무게를 초과하지 않으면서 보석의 가격 합이 최대가 되도록 해야 한다.

알고리즘 설명

1. 보석을 무게 기준으로 오름차순 정렬한다.
2. 가방을 무게 기준으로 오름차순 정렬한다.
3. 각 가방에 대해, 해당 가방에 담을 수 있는 모든 보석을 선택해 힙에 넣는다.
4. 힙에서 가장 비싼 보석을 선택해 결과에 더한다.

시간복잡도

1. 보석 정렬: O(Nlog⁡N)
2. 가방 정렬: O(Klog⁡K)
3. 각 가방에 대해 보석을 힙에 넣는 작업: 최악의 경우 모든 보석을 힙에 넣는 경우 O(Nlog⁡N)
4. 힙에서 가장 비싼 보석을 꺼내는 작업: 최악의 경우 모든 가방에 대해 힙에서 보석을 꺼내는 경우 O(Klog⁡N)

따라서, 전체 시간복잡도는 O(Nlog⁡N+Klog⁡K+Klog⁡N)로, 주로 Nlog⁡N와 Klog⁡N에 의해 결정된다.

from heapq import heappop, heappush  # 힙 관련 함수 임포트
import sys; input = sys.stdin.readline  # 빠른 입력 함수 설정

n, k = map(int, input().split())  # 보석 개수 n과 가방 개수 k 입력
jewels = [list(map(int, input().split())) for _ in range(n)]  # 각 보석의 무게와 가격 입력
jewels.sort()  # 보석을 무게 기준으로 정렬
bags = [int(input()) for _ in range(k)]  # 각 가방의 최대 무게 입력
bags.sort()  # 가방을 무게 기준으로 정렬
res = 0  # 훔칠 수 있는 보석 가격의 합의 최댓값 초기화
tmp = []  # 힙을 사용하여 가격이 높은 보석을 임시 저장할 리스트

for bag in bags:  # 각 가방에 대해
    while jewels and jewels[0][0] <= bag:  # 보석이 남아있고, 보석의 무게가 가방의 무게 이하인 경우
        heappush(tmp, -heappop(jewels)[1])  # 해당 보석의 가격을 힙에 추가 (최대힙을 사용하기 위해 음수로 저장)
    if tmp:  # 힙이 비어있지 않다면
        res -= heappop(tmp)  # 힙에서 가장 큰 값을 꺼내어 결과에 더함 (음수로 저장했으므로 다시 음수로 변환하여 더함)

print(res)  # 최종 결과 출력