[Python/파이썬] 프로그래머스 더 맵게

프로그래머스

코드

import heapq

def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)
    while True:
        # 가장 맵지 않은 음식의 스코빌 지수가 k가 되면 break
        if scoville[0] >= K:
            break
        # 모든 음식의 스코빌 지수를 K 이상으로 만들 수 없는 경우
        if len(scoville) == 1 and scoville[0] < K:
            return -1
        mapzzil = heapq.heappop(scoville)
        mapzzil2 = heapq.heappop(scoville)
        heapq.heappush(scoville, mapzzil+mapzzil2*2)
        answer += 1
    return answer

파이썬의 heapq 모듈을 이용하면 heap 자료구조를 손쉽게 사용할 수 있다. 파이썬의 heapq는 최소힙을 제공하므로 이를 이용하면 이번 문제를 쉽게 풀 수 있다.

1. 우선 scoville 배열을 heapify 함수를 이용하여 heap 자료 구조로 만들어준다. 최소힙이므로 항상 가장 맵지 않은 음식이 0번 인덱스에 있게 될 것이다.
2. 0번 인덱스의 값이 K 이상이 되면 문제의 조건을 만족하므로 반복문을 탈출한다. 이 때 주의할 점은 scoville 배열의 길이가 1이 된다면 음식이 1개 남은 것인데 이 음식의 스코빌 지수가 K를 넘지 못한다면 이 음식들로는 스코빌 지수가 K이상인 음식들을 결국 만들 수 없는 것이므로 1을 리턴해주도록 한다. 
   위의 2개의 if 문에 걸리지 않는다면 아직 음식을 계속 섞어봐야 하는 것이므로 scoville 힙에서 원소 2개를 뽑아 문제에서 제시한 수식대로 계산하여 scoville 힙에 heappush 함수를 이용하여 다시 넣어주고 answer의 값을 증가시켜준다.