[프로그래머스] 더 맵게 - python

📖 문제 링크

https://programmers.co.kr/learn/courses/30/lessons/42626

📃 문제 설명

힙을 이용해서 푸는 문제다.

K이상의 스코빌 지수가 될 때까지 2가지의 음식을 섞어 모든 음식이 K 이상의 스코빌 지수가 되도록 하고 횟수를 구하는 문제다.

파이썬에는 heapq이 있는데 이녀석의 능력은 push나 pop을 할 때마다 자동으로 정렬해주는 역할을 하고 있다.

📃 조건

• 모든 음식의 스코빌 지수를 K 이상으로 만들 수 없는 경우 -1 return

👨‍💻 문제 풀이

처음 풀이

def solution(scoville, K):
    count = 0
        
    while scoville[0] < K:
        change_num = scoville[0] + scoville[1] * 2
        scoville[0] = change_num
        del(scoville[1])
        scoville.sort()
        count += 1
        
    
    return count

할 때마다 sort와 del을 하는 바람에 시간 복잡도가 많이 올라가고 시간 초과가 떳다..

정답

import heapq as hq

def solution(scoville, K):
    heap = []

    for num in scoville:
        hq.heappush(heap, num)

    count = 0

    while heap[0] < K:
        try:
            hq.heappush(heap, hq.heappop(heap) + (hq.heappop(heap) * 2))
        except IndexError:
            return -1
        count += 1


    return count

그래서 heapq을 사용했는데
try부분에서 pop으로 제거하다가 없으면 IndexError가 뜨게 되고 그러면 -1로 return하도록 했다.

heappush나 heappop같은 경우 바로 적용되기 때문에 따로 변수를 지정해주지 않아도 된다.

이번에 풀면서 heapq을 처음 써보았는데 heapq에 대해 잘 알아봐야겠다.

2021.08.24