[백준] 1202. 보석 도둑 - 파이썬

[Gold II]

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

1202번: 보석 도둑

 

풀이 1

담을 수 있는 보석의 value를 최대로 해야 하므로, 일단은 우선순위 큐 (최대 힙)을 이용할 것이라 생각했음.

우선 생각나는 대로 먼저 구현해 보았고, line 18 ~ line 33 쪽에서 가방에 들어갈 수 있는 지 Brute-Force로 반복해가며 검사했기 때문에, 아마 TLE(시간초과)가 나올 것이라고 예상했다.

대충 계산해 봐도 Time Complexity가 가방 개수 K, 보석 개수 N일 때, O(K^2 logN) 정도기 때문에..

1 <= N, K <= 300,000 범위라.. N과 K가 거의 비슷하다고 치면.. O(N^2 logN)임. 될 리가 없지.

 

결과는 시간초과.

import heapq
N, K = map(int, input().split())

jewels = []
'''보석의 value가 첫 번째 요소가 되도록, value 기준 최대 힙으로 정렬.'''
for _ in range(N):
    a = tuple(map(int, input().split()))[::-1]
    jewels.append((-a[0], a[1]))

heapq.heapify(jewels)

bags = []
for _ in range(K):
    bags.append(int(input()))
bags.sort()

total_value = 0
for _ in range(N):
    '''보석을 최대 힙에서 하나씩 pop하며 확인'''
    jewel = heapq.heappop(jewels)
    value = -jewel[0]
    weight = jewel[1]

    '''가방 array는 무게순으로 정렬되어 있음'''
    for i in range(len(bags)):
        if weight > bags[i]:
            '''보석 무게가 가방 허용범위보다 크다면'''
            continue
        else:
            '''보석이 가방에 들어갈 수 있다면'''
            del bags[i]
            total_value += value
            break

print(total_value)

 

풀이 2

정말 당연하게도, value가 최대인 보석의 순서대로 뽑고, 그 때마다 그 보석을 담을 수 있는 가방 중, 허용되는 무게가 최소인 가방을 고르는 것이 최적해를 찾는 방법이라고 생각했다.

Ex. ((value, weight) 형태로 표현) 보석이 3개 : (99, 2), (65, 1), (23, 5) 이고,

가방이 2개 : (허용 무게: 10, 2) 일 때,

당연히 (99, 2)인 보석을 넣는 가방은 허용 무게가 2여야 하고, (65, 1)을 넣는 가방은 허용 무게가 10이여야 한다고 생각했다.

그러나,

힌트 : 두 번째 예제(위와 같은 예제)의 경우 첫 번째 보석을 두 번째 가방에, 세 번째 보석을 첫 번째 가방에 넣으면 된다.

를 보니, 최적해에 집중하는 것이 아닌 문제라고 생각했다.

최적해가 아니다? 그렇다면? Greedy.

와 같은 사고 순서로.. Greedy하게, 그냥 가장 허용 무게 범위가 가장 큰 가방부터 보석을 담는 식으로 진행해봤다.

물론 이렇게 하면 Time Complexity는 O(N logN)정도.

 

결과는 WA(틀렸습니다.)

아니 근데 이게 아닌 것 같긴 했어..

import heapq
N, K = map(int, input().split())

jewels = []
'''보석의 value가 첫 번째 요소가 되도록, value 기준 최대 힙으로 정렬.'''
for _ in range(N):
    a = tuple(map(int, input().split()))[::-1]
    jewels.append((-a[0], a[1]))

heapq.heapify(jewels)

bags = []
for _ in range(K):
    bags.append(int(input()))

'''가방 내림차순 정렬'''
bags.sort(key=lambda x: -x)

total_value = 0
for _ in range(N):
    '''보석을 최대 힙에서 하나씩 pop하며 확인'''
    jewel = heapq.heappop(jewels)
    value = -jewel[0]
    weight = jewel[1]

    '''가방 array는 무게 내림차순으로 정렬되어 있음. bags[0]은 가방들 중 항상 최대의 무게를 수용할 수 있음.'''
    if not bags:
        break
    if weight > bags[0]:
        '''보석 무게가 가방 허용범위보다 크다면'''
        continue
    else:
        '''보석이 가방에 들어갈 수 있다면'''
        del bags[0]
        total_value += value

print(total_value)

 

풀이 3

결국 풀이 1같이 가장 작은 가방에 가능한 한 제일 보석을 넣어야 된다는 건 맞는 것 같다.

그럼 풀이 1에서 조금 최적화를 시도해 보자.

 

풀이 1에서는 모든 보석을 하나씩 pop해가며 들어갈 수 있는 가방 중 가장 작은 가방을 택했다.

이를 순서를 바꿔, 가장 작은 가방부터, 넣을 수 있는 보석 중 가장 value가 높은 것을 고르는 식으로 해보자.

 

근데 이거 풀다 보니, 풀이 1의 line 18 ~ line 33 만 바꾸는 식으로, (즉, 보석의 value에 대한 최대 힙을 미리 만들어놓는 식으로) 풀이하면 결국 똑같이 O(N^2 logN)정도의 Time Complexity가 나오는 꼴이라.. 방향을 바꿔야겠다고 생각했다.

 

그렇게 생각한 것이 아래 풀이.

가장 작은 가방부터 시작하는 것은 같으나, 보석 또한 무게순으로 정렬한 후,

가방마다 들어갈 수 있는 모든 보석을 value 기준의 max heap에 넣는다.

그리고, 들어갈 수 있는 보석이 남아있다면, max heap의 맨 위에서 pop하여 (가장 value가 높은 것을 택하여) 가방에 넣는다.

이런 식으로 하면, 우선 가장 작은 가방부터 시작하므로, max heap에 들어가 있는 보석은 뒤에 나오는 모든 가방에 들어갈 수 있게 된다.

또한, 모든 보석에 대해서 한 번씩만 check하게 되기 때문에, (만약 보석이 한 번 들어갈 수 없는 무게라면, pass 하게 되며 다음 순회에 다시 한 번 다른 가방에 대해 검사하지만, 그렇게 해도 보석의 개수가 N개라 할 때 2N번 이상으로 if문을 지나지 않음.) Time Complexity가 약 O(NlogN) 정도로 줄어든다!!

 

일단 이 코드는 TLE(시간 초과)가 났는데. 아무리 생각해 봐도 시간 초과가 난 것이 이해가 안 돼 이것저것 막 바꿔봤음. 입출력 속도 문제인가 싶어 sys를 이용해 입출력 최적화 했는데도 TLE 남. 이유는 아래에서.

import heapq
N, K = map(int, input().split())

jewels = []
'''보석을 weight 오름차순으로 정렬.'''
for _ in range(N):
    jewels.append(tuple(map(int, input().split())))
jewels.sort()

bags = []
for _ in range(K):
    bags.append(int(input()))

'''가방 weight 오름차순 정렬'''
bags.sort(key=lambda x: -x)

total_value = 0

temp = []
'''모든 가방에 대해 순회'''
for bag in bags:
    '''남은 보석이 존재하고, 보석이 가방에 들어갈 수 있다면 계속 heap에 push.'''
    while jewels:
        if jewels[0][0] <= bag:
            '''temp heap에 넣을 때는, value가 큰 순의 max heap으로. 가능한 것 중 제일 큰 value를 가진 보석을 택해야 함.'''
            '''이 부분이 Greedy!'''
            heapq.heappush(temp, (-jewels[0][1], jewels[0][0]))
            del jewels[0]
        else:
            '''남아있는 보석 중 가장 가벼운 보석도 넣지 못한다면, 그 가방은 넣을 보석이 없음.. ^ㅁ^'''
            break

    '''넣을 보석이 남아 있다면, 제일 value가 큰 보석을 택해 넣음.'''
    '''기본적으로 제일 작은 가방->큰 가방 순으로 순회했기 때문에, **temp heap에 들어 있는 보석은 모두 가방에 들어갈 수 있는 상태!!**'''
    if temp:
        total_value += -heapq.heappop(temp)[0]

print(total_value)

 

위 코드가 시간초과였던 이유는,, line 28의 del jewels[0] 때문이였음.

del ~ 함수 자체가, index로 접근하면서, O(N)의 Time Complexity를 가지기 때문에,

전체 Time Complexity가 O(N^2)이 되어버리는.. 사태가 발생합니다.

 

따라서 del jewels[0] 만 heapq.heappop(jewels) 로 바꿔주는 식으로 해결하였음.

어짜피 heappop도 첫 번째 원소를 pop하니까.. ^ㅁ^

 

아래는 최종 코드.

import sys
import heapq
input = sys.stdin.readline
N, K = map(int, input().split())

jewels = []
'''보석을 weight 오름차순으로 정렬.'''
for _ in range(N):
    jewels.append(tuple(map(int, input().split())))
jewels.sort()

bags = []
for _ in range(K):
    bags.append(int(input()))

'''가방 weight 오름차순 정렬'''
bags.sort()

total_value = 0

temp = []
'''모든 가방에 대해 순회'''
for bag in bags:
    '''남은 보석이 존재하고, 보석이 가방에 들어갈 수 있다면 계속 heap에 push.'''
    while jewels:
        if jewels[0][0] <= bag:
            '''temp heap에 넣을 때는, value가 큰 순의 max heap으로. 가능한 것 중 제일 큰 value를 가진 보석을 택해야 함.'''
            '''이 부분이 Greedy!'''
            heapq.heappush(temp, (-jewels[0][1], jewels[0][0]))
            heapq.heappop(jewels)
        else:
            '''남아있는 보석 중 가장 가벼운 보석도 넣지 못한다면, 그 가방은 넣을 보석이 없음.. ^ㅁ^'''
            break

    '''넣을 보석이 남아 있다면, 제일 value가 큰 보석을 택해 넣음.'''
    '''기본적으로 제일 작은 가방->큰 가방 순으로 순회했기 때문에, **temp heap에 들어 있는 보석은 모두 가방에 들어갈 수 있는 상태!!**'''
    if temp:
        total_value += -heapq.heappop(temp)[0]

print(total_value)