백준 1202번. 보석 도둑(파이썬)

자세한 문제 내용은

아래 링크로 가주세요.

www.acmicpc.net/problem/1202

1202번: 보석 도둑

접근 방법

---

답을 구하는 법 자체는 어렵지 않은데, 시간 초과 때문에 애먹었던 문제입니다.

보석의 무게와 가격 M, V가 있고, 가방에 담을 수 있는 최대 C가 있습니다.

그럼 M이 C이하인 것 중에 가장 V가 큰 원소를 구해서 더하면 되는 문제라고 생각해서 접근해봤습니다.

N, K = stdin.readline().split()
N = int(N)
K = int(K)

일단 N과 K를 입력받아서 각각 int형으로 저장합니다.

for i in range(N):
    x, y = stdin.readline().split()
    x = int(x)
    y = int(y)
    li.append([x, y])
for i in range(K)
	x = int(stdin.readline().split()
    C.append(x)
MV.sort()
C.sort()

그리고 이렇게 보석 정보 리스트 li과 가방에 들어갈 수 있는 무게 리스트 C로 나눈 뒤에 두 리스트를 오름차순으로 정렬합니다.

for i in C:
    for j in range(len(MV)):
        if MV[j][0] <= i:
            res += MV[j][1]
            MV.remove(MV[j])
            break

그런 뒤 중첩 for문을 사용해서 정답을 구해냇었는데, 답 자체는 맞는 것 같은데 제출해보니 시간 초과가 발생했습니다.

그래서 다음으로는 PriorityQueue(우선순위 큐)를 통한 방법을 사용해봤습니다.

우선순위 큐는 값을 넣을 때는 보통 큐와 다름없지만, 값을 꺼낼 때에는 가장 작은 값부터 꺼냅니다.

for i in range(N):
    x, y = stdin.readline().split()
    x = int(x)
    y = int(y)
    li.append([x, y])
for i in range(K):
    li.append([int(stdin.readline()), 1000001])
li.sort()

일단 두 리스트를 하나도 합쳐서 위와 같이 수정합니다.

먼저 보석의 정보(M, V)를 입력받고, 가방의 무게(C)와 보석 가격의 최댓값+1을 추가시켰습니다.

그리고 이 리스트를 정렬시켜줍니다.

 

예를 들어 보석 정보를 [1, 3000], [1, 6000], [1, 1000]과 같이 입력하고, 가방 무게를 [1, 1000001]과 같이 입력했다고 합시다.

그러면 보석들은 같은 무게에서 가격이 오름차순으로 정렬되고, 가장 뒤에 가방 정보가 위치하게 됩니다.

([1, 1000], [1, 3000], [1, 6000], [1, 1000001]과 같이)

que = queue.PriorityQueue(N)
for i in li:
    if i[1] != 1000001:
        que.put(-i[1])
    else:
        if not que.empty():
            res += (-que.get())

이제 다음으로 PriorityQueue를 최대 N의 크기만큼 생성시켜줍니다.

그리고 for문을 사용해서 li리스트의 요소를 하나씩 가져오면서 조건을 검사합니다.

i[1]이 1000001이 아니면, 즉 보석일 경우 queue에 해당 보석의 가격을 queue에 put 해줍니다.

 

여기서 가격에 -를 붙이는 이유는 우선순위 큐에서는 가장 작은 값을 꺼내는데, 우리가 구해야 하는 것은

가장 비싼 가격의 보석이기 때문에 -를 붙여서 값이 가장 큰 보석이 우선적으로 나올 수 있게 하기 위함입니다.

 

요소가 보석이 아닐 경우, 즉 가방의 정보일 경우 먼저 queue가 비어있지 않은지 검사합니다.

queue가 비어있지 않다면 queue에서 가장 작은 값(가장 비싼 보석의 가격)을 get 해서 다시 부호를 바꿔 res에 더해줍니다.

그림으로 표현하자면 이렇게 되는데, 보석 가격을 큐에 쌓아두다가 가방 정보가 들어올 경우

현재 쌓여있는 보석 중에서 가장 작은(가장 비싼) 보석의 가격을 큐에서 빼내 결과(res)에 더해줍니다.

전체 소스코드는 아래에 있습니다.

import queue
from sys import stdin
N, K = stdin.readline().split()
N = int(N)
K = int(K)
li = []
res = 0
for i in range(N):
    x, y = stdin.readline().split()
    x = int(x)
    y = int(y)
    li.append([x, y])
for i in range(K):
    li.append([int(stdin.readline()), 1000001])
li.sort()
que = queue.PriorityQueue(N)
for i in li:
    if i[1] != 1000001:
        que.put(-i[1])
    else:
        if not que.empty():
            res += (-que.get())
print(res)