프로그래밍/Algorithm

백준 1202 보석도둑 파이썬

모지사바하 2021. 3. 10. 11:28

www.acmicpc.net/problem/1202

 

1202번: 보석 도둑

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000) 다음 N개 줄에는 각 보석의 정보 Mi와 Vi가 주어진다. (0 ≤ Mi, Vi ≤ 1,000,000) 다음 K개 줄에는 가방에 담을 수 있는 최대 무게 Ci가 주어진다. (1 ≤ Ci

www.acmicpc.net

import heapq
N, K = map(int, input().split())
gem = []
for i in range(N):
    weight, value = map(int, input().split())
    heapq.heappush(gem, [weight, value])

bag = []
for i in range(K):
    heapq.heappush(bag, int(input()))
    

in_bag = []
total_value = 0

for _ in range(K):
    bag_capacity = heapq.heappop(bag)

    while gem and bag_capacity >= gem[0][0]:
        [weight, value] = heapq.heappop(gem)
        heapq.heappush(in_bag, -value)
    
    if in_bag:
        total_value-=heapq.heappop(in_bag)
    elif not gem:
        break
print(total_value)        

 

혼자 풀어보려했으나,

시간초과, 메모리초과가 계속 나와서 결국 다른분 블로그를 참고했다.

아직 알고리즘에 자료구조 쓰는걸 떠올리지 못하는듯하다.

최소힙, 최대힙을 이용하니 굉장히 심플하고 간결하게 구현이 가능하다.

매번 잘 못풀어서 다른 블로그를 참고하지만 그래도 확실히 이해하고 넘어가도록 하자.

 

1. 우선 보석을 무게순으로 최소힙으로 만든다.

2. 가방도 최소힙으로 만든다. (가방은 굳이 최소힙으로 만들 필요는 없고 리스트 오름차순 정렬만 해도 됨)

3. 가방의 갯수만큼 순회하면서

4. 현재 가방의 수용 무게를 선택한다. 이때 가방은 최소힙이기 때문에 최소 무게가 선택된다.

5. 훔칠 보석이 있고, 보석의 무게가 가방의 수용가능 무게 이하라면 순회하면서

(첫번째 순회라고 했을때 현재 선택된 가방은 최소 수용가능한 무게다. 그리고 보석역시 최소힙이기 때문에  최소의 무게가 선택되는데 이 무게가 가방의 수용가능 무게보다 무겁다면 현재 선택된 가방으로는 훔칠 수 있는 보석이 없으므로 고려대상이 없다.)

6. 현재 가방에 담을 수 있는 무게의 모든 보석을 일단 담는데 담을때 그냥 담는게 아니고 보석가격을 최대힙으로 담는다. 파이썬은 기본적으로 최소힙이기 때문에 음수로 담은후 꺼낼때 - 를 곱해주는 방법으로 최대힙을 이용할 수 있다.

최대합으로 담는 이유는 일단 담을수 있는 보석은 다 담았지만 그 중 가장 비싼 보석을 훔쳐야하기때문에 보석가격을 최대힙으로 담은 것이다.

7. 현재 가방에 담을 수 있는 보석을 다 담았다면, 즉 while 순회가 끝났다면

8. 가방에 보석이 담겨있다면 가방에서 첫번째 요소를 꺼내서 결과에서 빼준다. 빼주는 이유는 가방에 보석가격을 최대힙으로 담기위해 보석의 가격이 음수로 담겨있기때문이다. 음수를 빼주면 양수가 된다. 즉 결과에 누적시키는 것과 같은 결과가 나온다.

9. 가방에 보석도 없고 훔칠보석도 남아있지 않다면 루프 종료.