[프로그래머스/Lv2] 더 맵게

문제

매운 것을 좋아하는 Leo는 모든 음식의 스코빌 지수를 K 이상으로 만들고 싶다.

모든 음식의 스코빌 지수를 K 이상으로 만들기 위해 Leo는 스코빌 지수가 가장 낮은 두 개의 음식을 아래와 같이 특별한 방법으로 섞어 새로운 음식을 만든다.

 

섞은 음식의 스코빌 지수 = 가장 맵지 않은 음식의 스코빌 지수 + (두 번째로 맵지 않은 음식의 스코빌 지수 * 2)

 

Leo는 모든 음식의 스코빌 지수가 K 이상이 될 때까지 반복하여 섞는다.

Leo가 가진 음식의 스코빌 지수가 담긴 배열 scoville과 원하는 스코빌 지수 K가 주어질 때, 모든 음식의 스코빌 지수를 K 이상으로 만들기 위해 섞어야 하는 최소 횟수를 반환하는 함수를 만들어야 한다.

 

 

제한사항

1. scoville의 길이는 2 이상 1,000,000 이하이다.

2. K는 0 이상 1,000,000,000 이하이다.

3. scoville의 원소는 각각 0 이상 1,000,000 이하이다.

4. 모든 음식의 스코빌 지수를 K 이상으로 만들 수 없는 경우에는 -1을 반환한다.

 

 

코드

import heapq

def solution(scoville, K):
    i = 0
    check = True
    heapq.heapify(scoville)
    while check:
        try:
            heapq.heappush(scoville, heapq.heappop(scoville) + (heapq.heappop(scoville) * 2))
            if scoville[0] >= K:
                check = False
            i += 1
        except:
            i = -1
            break
    
    return i

 

 

설명

단순 배열로 정렬과 pop 등을 사용해서 문제를 풀 순 있는데

효율성 테스트에서 전부 시간초과된다.

 

프로그래머스에서 이 문제의 카테고리는 힙(Heap)으로 분류되어 있고, 실제로 힙을 사용해서 풀면 간단하게 통과된다.

파이썬에는 heapq 알고리즘을 제공하는 모듈이 내부에 존재하기 때문에 힙을 구현할 수 있다.

 

i = 0
check = True

i는 반복 횟수를 저장할 변수이고 check는 while 반복문의 트리거 변수이다.

 

heapq.heapify(scoville)

입력으로 받은 scoville 배열을 heapq모듈의 heapify 함수를 이용해 힙 자료형으로 변환한다.

파이썬의 heapq모듈은 배열을 최소 힙의 형태로 정렬해준다.

 

while check:
    try:
        heapq.heappush(scoville, heapq.heappop(scoville) + (heapq.heappop(scoville) * 2))
        if scoville[0] >= K:
            check = False
        i += 1
    except:
        i = -1
        break

check가 True인 상태라면 계속해서 반복한다.

scoville 힙에서 가장 작은 수를 pop하고, 그 후에 같은 동작을 해서 두 번째로 작은 수를 pop한다.

이 두 수로 문제에 나온대로 계산하여 다시 scoville 힙에 push한다.

 

그 후에 scoville의 가장 첫 번째 원소가 K보다 크거나 같다면 check를 False로 바꾼다(=반복문을 종료).

i를 1 더해서 반복 횟수를 기록한다.

 

만약 에러가 발생한다면(K보다 큰 수를 만들 수 없다면) i를 -1로 초기화하고 반복문을 강제 종료한다.

 

return i

구해진 i를 반환한다.