[Programmers / Python / 힙] - 더 맵게

Young·2021년 5월 30일
0

출처 https://programmers.co.kr/learn/courses/30/lessons/42626?language=python3

모든 음식의 스코빌 지수를 K 이상으로 만들려 한다.
그러기 위해 스코빌 지수가 가장 낮은 두 개의 음식을 아래와 같이 특별한 방법으로 섞어 새로운 음식을 만든다.
Leo는 모든 음식의 스코빌 지수가 K 이상이 될 때까지 반복하여 섞습니다.

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

scoville : Leo가 가진 음식의 스코빌 지수를 담은 배열
K : 원하는 스코빌 지수
solution 함수 : 모든 음식의 스코빌 지수를 K 이상으로 만들기 위해 섞어야 하는 최소 횟수 return

나의 풀이

  1. scoville을 binary min heap으로 변환
  2. min heap의 root 가 K가 될 때 까지 음식을 섞는 행위를 한다.
    (음식 2개를 꺼내 섞은 후 그 1개를 다시 넣음)
import heapq

def solution(scoville, K):
    answer = 0
    
    heapq.heapify(scoville)

    while scoville[0] < K:
        heapq.heappush(scoville, heapq.heappop(scoville) + heapq.heappop(scoville) * 2)
        answer += 1

    return answer

몇 가지 테스트 케이스에서 런타임 오류가 발생했다. 만약 pop할 원소가 없을 때 pop을 시킨다면 런타임 오류가 발생할텐데 내가 이 부분을 처리해주지 않은 것을 알게 되었고 문제를 다시 보니 조건에도 이렇게 나와있었다.

그래서 heap의 원소 갯수가 2보다 작은 경우 pop 하지 않고 바로 -1을 return 해주도록 코드를 수정했다.

import heapq

def solution(scoville, K):
    answer = 0
    
    heapq.heapify(scoville)

    while scoville[0] < K:
        if len(scoville) < 2:
            return -1
        heapq.heappush(scoville, heapq.heappop(scoville) + heapq.heappop(scoville) * 2)
        answer += 1

    return answer

profile
👩🏻‍💻

0개의 댓글