프로그래머스 | 더 맵게 [파이썬 python]

프로그래머스

 

코드

import heapq

def solution(scoville, K):
    answer = 0 # 섞은 횟수
    heapq.heapify(scoville) # 최소 heap
    
    # 새로운 스코빌 지수 계산이 가능할 때(len >= 2)
    while len(scoville) >= 2:
        min1 = heapq.heappop(scoville)
        # 최소값이 K 이상이면 조건을 만족하므로 반복문 종료
        if min1 >= K:
            return answer
        # 최소값이 K 미만이면 새로운 스코빌 지수 계산
        else: 
            min2 = heapq.heappop(scoville)
            new = min1 + (min2 * 2) 
            heapq.heappush(scoville, new)
            answer += 1
    
    # 마지막으로 하나의 음식만 남았을 때
    if scoville[0] >= K:
        return answer
    else:
        return -1

heapq를 사용해야만 효율성 테스트를 통과할 수 있는 문제이다. 

처음에는 단순히 주어진 리스트만으로 풀었지만 효율성에서 시간 초과가 발생해 자동으로 정렬을 하는 heap을 사용하였다.

 

먼저, heapify를 사용해 주어진 리스트를 heap으로 만든다. 

heapq.heapify(scoville) # 최소 heap

 

문제에서 가장 스코빌 지수가 낮은 두 값을 가지고 새로운 스코빌 지수를 만든다고 했다.

그러므로 새로운 스코빌 지수 계산을 하기 위해서는 최소 2개의 원소가 존재해야 하므로 len을 사용해 원소가 2개 이상인 경우에 스코빌 지수 계산을 반복한다.

# 새로운 스코빌 지수 계산이 가능할 때(len >= 2)
while len(scoville) >= 2:
    min1 = heapq.heappop(scoville)
    # 최소값이 K 이상이면 조건을 만족하므로 반복문 종료
    if min1 >= K:
        return answer
    # 최소값이 K 미만이면 새로운 스코빌 지수 계산
    else: 
        min2 = heapq.heappop(scoville)
        new = min1 + (min2 * 2) 
        heapq.heappush(scoville, new)
        answer += 1

스코빌 지수 계산을 위해 먼저 heappop을 통해 최솟값을 가져와 min1로 저장하고, 최솟값 min1이 K 이상인 경우에는 조건을 만족하므로 answer를 리턴하여 반복문을 종료한다.

그렇지 않으면 새로운 스코빌 지수 계산을 위해 다음 최솟값을 가져와 min2로 저장하고 새로운 스코빌 지수를 scoville에 추가한다.

 

scoville의 길이가 2 미만 즉, 1이 되면 while문을 탈출한다.

이때, 남은 원소 하나가 조건을 충족하는지 확인하여 충족하면 answer를 그렇지 않으면 -1을 리턴한다.

# 마지막으로 하나의 음식만 남았을 때
if scoville[0] >= K:
    return answer
else:
    return -1