scoville | K | return |
---|---|---|
[1, 2, 3, 9, 10, 12] | 7 | 2 |
가장 작은 수와 두번째로 작은 수를 계산한 후 이를 포함한 리스트에서 또 가장 작은 수와 두번째로 작은 수를 찾아내는 것 반복.. 그때마다 sort를 해주는 것은 너무 번거로우니 최소힙을 제공하는 heapq를 사용하자
import heapq
def solution(scoville, K):
heapq.heapify(scoville)
ans = 0
while len(scoville) >= 2:
if scoville[0] >= K:
return ans
heapq.heappush(scoville, heapq.heappop(scoville) + 2*(heapq.heappop(scoville)))
ans += 1
if scoville[0] >= K:
return ans
return -1
마지막에 scoville원소 개수가 2개이고 이를 계산한 결과인 원소는 K이상일 수도 있다. 따라서 while문이 끝난 후에 한번 더 조건문으로 판별해줘야한다.