[힙] 프로그래머스 - 더 맵게

문제 출처 : https://programmers.co.kr/learn/courses/30/lessons/42626

 

문제 설명 : 

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

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

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

 

문제 풀이 :

import heapq
def solution(scoville, K):
    count = 0
    heapq.heapify(scoville)
    try:
        while scoville[0] < K :
            a = heapq.heappop(scoville)
            b = heapq.heappop(scoville)
            heapq.heappush(scoville, a + b*2)
            count += 1
    except IndexError:
        count = -1
    return count

최소 힙 정렬을 이용하려고 heapq 모듈을 불러왔다. 주어진 scoville을 heapify를 이용해서 최소 힙으로 만들고, 0번째 원소가 K보다 작다면 0번째 원소와 1번째 원소를 pop해서 섞은 후에 다시 힙에 넣어준다. 이 과정을 반복하다보면 0번째 원소가 K보다 커져서 반복문이 종료되거나 남은 원소가 한 개여서 결국 K를 넘는 음식을 만들지 못하게 된다. 전자의 경우는 섞은 횟수를, 후자의 경우는 IndexError 예외 처리로 -1을 return 해주면 된다.