[프로그래머스] 힙(heap) - 더 맵게 (Python)

코딩테스트 연습 - 더 맵게

내 풀이

import heapq
def get_sco(d1, d2) : 
    return d1+d2*2

def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)
    while True :
        if len(scoville) == 1 and scoville[0] < K :
            return -1
        first = heapq.heappop(scoville)
        second = heapq.heappop(scoville)
        heapq.heappush(scoville, get_sco(first, second))
        answer += 1
        if scoville[0] >= K :
            return answer

 스코빌 지수를 위해서는 모든 원소가 K보다 크거나 같아야 한다. 그 작업을 위해서는 sort를 계속해주는 것이 필요한데 그러기엔 효율성에서 계속 걸리고 heap을 잘 쓰는 것이 관건이었다.

heapify를 통해 list을 heap형태로 바꾸어주고, 앞의 두 원소를 리스트에 집어 넣어줌으로써 mix_sco를 통해 스코빌 지수를 구해주었다. 작업을 실행할 때마다 answer를 1씩 증가시켜주며 몇 번 작업했는지 체크해주었고, 모든 작업 후에도 scovile이 K보다 작다면 -1을 return 해줌

다른 풀이

조금 더 예쁘게 코드를 고쳐볼 수 있었다.

import heapq
def mix_sco(dt1, dt2) :
    return dt1+dt2*2

def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)
    while scoville[0] < K :
        first = heapq.heappop(scoville)
        second = heapq.heappop(scoville)
        heapq.heappush(scoville, mix_sco(first, second))
        answer += 1
        if len(scoville)==1 and scoville[0] < K :
            return -1
    return answer

어차피 heapq의 형태이기 때문에 정렬은 되어 있었고, heappop한 것을 그대로 값으로 할당해줘서 코드를 조금 더 짧게 구현할 수 있었다.