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

문제 요약

• 가방에 최대한 가격이 높은 보석을 담았을 때 얻을 수 있는 보석 가격의 합을 구하는 문제이다.

• 한 가방에는 오직 하나의 보석만 들어갈 수 있다.

• N(보석의 수), K(가방의 수)의 최대 값이 각각 300,000으로, O(N * K)의 이중 for loop만 돌아도 시간 초과가 발생할 수 있다.

풀이

• 우선순위 큐를 통해 O((N+K)*logN)의 시간복잡도로 문제를 해결할 수 있다.

우선순위 큐(Priority Queue)

• 우선순위 큐란 말 그대로 우선순위가 있는 큐를 말한다.
• pop 연산을 수행했을 때 첫번째나 마지막에 있는 값을 빼는 것이 아니라, 우선순위가 가장 높은 값을 뺄 수 있는 자료구조를 말한다.
• 파이썬에서는 이를 힙(heap)을 통해 구현할 수 있다.

풀이 과정

1. 보석과 가방을 입력받은 뒤 각각 무거운 순으로 정렬하기

# 무게가 무거운 순으로 정렬
jewels = [tuple(map(int, input().split())) for _ in range(N)]
jewels.sort(key = lambda x : x[0], reverse = True)
# 무거운 가방 순으로 정렬
bags = [int(input()) for _ in range(K)]
bags.sort(reverse = True)

• 무거운 순으로 정렬하는 이유는 pop 연산을 했을 때 가벼운 것부터 순차적으로 꺼내기 위함이다.

2. 우선순위 큐 선언하기 (최대 힙)

# 우선순위 큐 (max heapque)
hq = []

• 어짜피 나중에 heappush 함수를 통해 힙 자료구조가 될 것이므로 그냥 빈 리스트를 선언하면 된다.

3. 가장 가벼운 가방부터 꺼내는 반복문 수행

while bags:
    # 가장 가벼운 가방을 꺼냄
    bag = bags.pop()

• while 문을 통해 가장 가벼운 가방을 꺼낸다.

4. 가장 가벼운 보석부터 꺼내어 가방에 넣을 수 있다면 우선순위 큐에 넣기

while jewels:
        # 현재 가장 가벼운 보석을 꺼냄
        weight, value = jewels.pop()
        if bag >= weight:
            # 가방에 넣을 수 있다면 최대힙에 추가
            heapq.heappush(hq, -value)
        else:
            # 넣을 수 없다면 다시 보석 리스트에 추가하고 탐색 종료
            jewels.append((weight, value))
            break

• 가벼운 순으로 보석을 꺼내 가방에 담을 수 있는 무게보다 작다면 우선순위 큐(힙)에 추가한다.
  • 파이썬의 힙은 기본적으로 최소힙이 되도록 작동하므로 음수 값을 heappush하면 절대값이 큰 것부터 빼오는 최대힙을 만들 수 있다.
• 보석을 가방에 담을 수 없다면 다시 리스트에 넣어준다.
• 이렇게 하면 heappop 연산을 수행했을 때, 가장 높은 가격의 음수 값이 반환된다. 즉 높은 가격이 높은 우선순위가 되는 우선순위 큐를 구현한 것이다.

5. 가방에 넣을 수 있는 가장 높은 가격(가장 높은 우선순위)의 보석을 정답 변수에 더하기

# 최대힙에서 heappop 한 값을 답변에 더함
if hq:
    answer -= heapq.heappop(hq)

• 이제 heappop된 값을 정답 변수에 더해주면 된다.
• 힙에 음수 값을 넣어줬으므로 빼기 연산을 해주어야 한다.
• 가방에 넣을 수 있는 보석이 없는 경우도 있으므로 반드시 힙에 값이 들어있는지 체크해주어야 한다. (빈 힙에서 pop하면 IndexError 발생)

예제

N=3 K=2
jewels = [(1, 65), (5, 23), (2, 99)]
bags = [10, 2]
answer = 0
hq = []

1. 보석과 가방을 각각 무거운 순으로 정렬한다.

jewels = [(5, 23), (2, 99), (1, 65)]
bags = [10, 2]
answer = 0
hq = []

2. 가장 가벼운 가방을 선택한다.

jewels = [(5, 23), (2, 99), (1, 65)]
bags = [10]
answer = 0
bag = 2
hq = []

3. 가장 가벼운 보석을 뽑는다. 보석의 무게는 1이므로 가방에 넣을 수 있다. 해당 보석의 가격을 우선순위 큐(힙)에 추가한다.

jewels = [(5, 23), (2, 99)]
bags = [10]
answer = 0
bag = 2
hq = [-65]

4. 다음 보석의 무게는 2로 가방에 넣을 수 있다. 마찬가지로 추가한다.

jewels = [(5, 23)]
bags = [10]
answer = 0
bag = 2
hq = [-99, -65] # 편의상 정렬된 리스트 형태로 표현

5. 다음 보석의 무게는 5이므로 가방에 넣을 수 없기 때문에 탐색을 종료한다. 이후 우선순위 큐에서 값을 꺼내어 정답 변수에 더해준다 (빼기 연산).

jewels = [(5, 23)]
bags = [10]
answer = 99
bag = 2
hq = [-65] # 편의상 정렬된 리스트 형태로 표현

6. 다음으로 무거운 가방을 선택한다.

jewels = [(5, 23)]
bags = []
answer = 99
bag = 10
hq = [-65] # 편의상 정렬된 리스트 형태로 표현

7. 마지막 보석을 꺼내고 무게를 비교해 우선순위 큐에 추가한다. 더이상 보석이 없으므로 탐색을 종료한다.

jewels = []
bags = []
answer = 99
bag = 10
hq = [-65, -23] # 편의상 정렬된 리스트 형태로 표현

8. 우선순위 큐에서 값을 꺼내 정답 변수에 더해준다.

jewels = []
bags = []
answer = 164
bag = 10
hq = [-23] # 편의상 정렬된 리스트 형태로 표현

9. 남은 가방이 없으므로 정답 변수를 출력하고 프로그램을 종료한다.

코드

import sys
import heapq

input = sys.stdin.readline

N, K = map(int, input().split())
jewels = []
bags = []
length = N
answer = 0

# 무게가 무거운 순으로 정렬
jewels = [tuple(map(int, input().split())) for _ in range(N)]
jewels.sort(key = lambda x : x[0], reverse = True)
# 무거운 가방 순으로 정렬
bags = [int(input()) for _ in range(K)]
bags.sort(reverse = True)
# 우선순위 큐 (max heapque)
hq = []

while bags:
    # 가장 가벼운 가방을 꺼냄
    bag = bags.pop()
    while jewels:
        # 현재 가장 가벼운 보석을 꺼냄
        weight, value = jewels.pop()
        if bag >= weight:
            # 가방에 넣을 수 있다면 최대힙에 추가
            heapq.heappush(hq, -value)
        else:
            # 넣을 수 없다면 다시 보석 리스트에 추가하고 탐색 종료
            jewels.append((weight, value))
            break
    # 최대힙에서 heappop 한 값을 답변에 더함
    if hq:
        answer -= heapq.heappop(hq)
print(answer)

코멘트

우선순위 큐에 대한 개념을 몰래 이분 탐색으로 풀려다가 시간초과로 한창 끙끙 댔다.

힙을 이용해 큐에서 꺼내는 순서를 원하는 대로 커스터마이징할 수 있다는 것을 잘 기억해두어야 겠다.

힙 구조에서는 pop, push 연산 모두 시간복잡도 O(log n)이라는 것도 기억해두자.