[파이썬/알고리즘] 프로그래머스 - 더 맵게 (Heap)

[문제]

매운 것을 좋아하는 Leo는 모든 음식의 스코빌 지수를 K 이상으로 만들고 싶습니다. 모든 음식의 스코빌 지수를 K 이상으로 만들기 위해 Leo는 스코빌 지수가 가장 낮은 두 개의 음식을 아래와 같이 특별한 방법으로 섞어 새로운 음식을 만듭니다.

섞은 음식의 스코빌 지수 = 가장 맵지 않은 음식의 스코빌 지수 + (두 번째로 맵지 않은 음식의 스코빌 지수 * 2)

Leo는 모든 음식의 스코빌 지수가 K 이상이 될 때까지 반복하여 섞습니다.
Leo가 가진 음식의 스코빌 지수를 담은 배열 scoville과 원하는 스코빌 지수 K가 주어질 때, 모든 음식의 스코빌 지수를 K 이상으로 만들기 위해 섞어야 하는 최소 횟수를 return 하도록 solution 함수를 작성해주세요.

제한 사항
scoville의 길이는 2 이상 1,000,000 이하입니다.
K는 0 이상 1,000,000,000 이하입니다.
scoville의 원소는 각각 0 이상 1,000,000 이하입니다.
모든 음식의 스코빌 지수를 K 이상으로 만들 수 없는 경우에는 -1을 return 합니다.

 

[풀이]

import heapq

def solution(scoville, K):
    heapq.heapify(scoville)
    count = 0
    
    while scoville[0] < K:
        if len(scoville) < 2:
            return -1
        
        min_first = heapq.heappop(scoville)
        min_second = heapq.heappop(scoville)
        
        new_scoville = min_first + (min_second * 2)
        
        heapq.heappush(scoville, new_scoville)
        
        count += 1
        
    return count

 

Heap 문제이다.

파이썬에서 import heapq를 통해서 쉽게 코드를 만들 수 있다.

 

heapq.heapify(scoville)

이 코드를 통해서 힙 구조로 변환할 수 있다.

예를 들어 [9, 3, 2, 5, 7] 이 있다고 할 때 힙 구조로 변환하게 되면 [2, 9, 3, 5, 7] 로 변하게 된다.

여기서 변환 될 때 볼 점은 가장 작은 원소가 맨 앞으로 배치가 되었다는 것이다. 하지만 그 뒤의 원소는 가장 작은 원소가 아니라는 것을 볼 수 있다. 따라서 힙 구조로 변환하면 가장 작은 원소가 맨 앞으로 오지만 그 뒤로도 작은 원소가 오는 것이 아니다.

 

while scoville[0] < K:

위에서 힙의 구조를 설명한 것처럼 항상 맨 앞에 있는 원소는 최솟값을 가지게 된다.
따라서 그 특성을 이용하여 가장 작은 값이 K 보다 크게 되면 루프를 멈추게 된다.

문제에서 모든 음식의 스코빌 지수를 K 이상으로 만들고 싶다고 했기 때문에 힙의 첫 번째 원소가 K보다 크다면 뒤에 값은 당연히 크다는 걸 알 수 있다.

 

if len(scoville) < 2:
     return -1

만약 scoville가 2개보다 적다면 만들 수 없으므로 -1을 반환한다.

 

 

min_first = heapq.heappop(scoville)

min_second = heapq.heappop(scoville)

첫 번째로 작은 값을 heappop을 통해서 빼고, 두 번째로 작은 값을 heappop을 통해서 빼게 된다.

여기서 의문이 드는 게, 위에서 최솟값은 정렬되었지만 그 뒤에 값은 정렬이 안되있는데 이렇게 빼는 게 맞냐 라고 할 수도 있다. 하지만 heap 특징인 가장 작은 원소가 맨 앞으로 오기 때문에 정렬이 안되있더라도 상관없다.

예를 들어 heap 구조의 [2, 9, 3, 5, 7]이 있다고 하면 첫 번째로 작은 값인 2를 min_first로 빼주게 되면 [3, 9, 5, 7] 이 남는다.

즉 haep의 특징에 따라서 3이 맨 앞으로 오게 되었고 여기서 또 heappop을 통해서 3이 min_second로 나오게 된다.

 

new_scoville = min_first + (min_second * 2)

그리고 문제에 나온 공식대로 계산해주게 된다.

 

heapq.heappush(scoville, new_scoville)

heappush를 통해서 기존 scoville에 new_scoville 값을 넣는다.

여기서도 heap의 특징을 유지하기 때문에 new_scoville가 scoville에서 최솟값이라면 맨 앞으로 가게 된다.

 

count += 1

루프가 한 번씩 돌때마다 1 씩 증가시켜서 문제에서 구하고자 하는 값을 구해준다.

 

return count

문제에서 구하고자 하는 값을 반환한다.