[프로그래머스] level2 더 맵게 - 파이썬(Python)

문제

코딩테스트 연습 - 더 맵게

 

코드

import heapq

def solution(scoville, K):
    heapq.heapify(scoville)
    first  = heapq.heappop(scoville)
    length = len(scoville)
    if(first >= K):
        return 0
    elif(length == 1 and scoville[0] < K):
        return 1
    else:
        for i in range(1, length):
            second = heapq.heappop(scoville)
            first  = heapq.heappushpop(scoville, first + second * 2)
            if first >= K: 
                return i
        return -1

 

코드풀이

테스트 16만 실패해서 엄청 고생했다... scoville = [1,2] , K = 3 넣어보세요. 1 나와야 합니다.

이 문제는 heap 에 대해 알고 있다면 쉽게 풀 수 있다.

 

1. heap.heapify() : heap 배열 변환

2. 가장 앞에 있는 최소값을 first 에 넣는다.

3. 현 scoville 의 길이를 length 에 넣는다.

4. 최소값이 이미 K 이상이라면 더 볼 필요도 없이 0 반환.

5. 만약 length 가 1이고 나머지 값이 K 보다 작다면 바로 1 반환

이외

 

heappushpop() : push 하고자 하는 값을 넣고 그 후 제일 작은 최소값을 pop한다.

scoville의 나머지 길이만큼 반복문을 돌려서 한번씩 최소값이 K 보다 큰 지 확인한다.

크면 섞은 횟수 i 반환. 다 돌렸는데도 안된다면 -1 반환. 

 

끝.