[Python] 더맵게

in Algorithm on 문제풀이, Heap


Programmers : 더맵게

뮨제 설명

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

더맵게1

Leo는 모든 음식의 스코빌 지수가 K 이상이 될 때까지 반복하여 섞습니다. Leo가 가진 음식의 스코빌 지수를 담은 배열 scoville과 원하는 스코빌 지수 K가 주어질 때, 모든 음식의 스코빌 지수를 K 이상으로 만들기 위해 섞어야 하는 최소 횟수를 return 하도록 solution 함수를 작성해주세요.

제한 사항

  • scoville의 길이는 1 이상 1,000,000 이하입니다.
  • K는 0 이상 1,000,000,000 이하입니다.
  • scoville의 원소는 각각 0 이상 1,000,000 이하입니다.
  • 모든 음식의 스코빌 지수를 K 이상으로 만들 수 없는 경우에는 -1을 return 합니다.

입출력 예

더맵게2

입출력 예 설명

  1. 스코빌 지수가 1인 음식과 2인 음식을 섞으면 음식의 스코빌 지수가 아래와 같이 됩니다. 새로운 음식의 스코빌 지수 = 1 + (2 * 2) = 5 가진 음식의 스코빌 지수 = [5, 3, 9, 10, 12]

  2. 스코빌 지수가 3인 음식과 5인 음식을 섞으면 음식의 스코빌 지수가 아래와 같이 됩니다. 새로운 음식의 스코빌 지수 = 3 + (5 * 2) = 13 가진 음식의 스코빌 지수 = [13, 9, 10, 12]

모든 음식의 스코빌 지수가 7 이상이 되었고 이때 섞은 횟수는 2회입니다.

문제에 들어가기 앞서 파이썬에서의 heap 모듈을 정리해보자.

heapify

heapq 모듈의 heapify 함수는 주어진 리스트를 힙 정렬한다. 이떄의 Time Complexity는 O(nlogn)입니다.

import heapq
my_list = [13, 2, 1, 5, 10]
heapq.heapify(my_list)

# 오름차순으로 정렬된다.
my_list # [1, 2, 13, 5, 10]

heappop

heap 모듈의 heappop 함수는 힙 정렬된 리스트에서

  1. 가장 작은 원소를 빼내고
  2. 나머지 원소가 힙을 유지하도록 정리

이 함수를 사용할 때에는 주어진 리스트가 힙 정렬되어있는지 반드시 확인. 정렬되지 않은 리스트에 heappop을 사용하면 이상한 결과가 나온다.

# heappop 예시 1
import heapq
my_list = [13, 2, 1, 5, 10]
heapq.heapify(mylist)

# 가장 작은 원소인 1이 리턴. mylist의 길이가 하나 줄어든다.
ret_val = heapq.heappop(my_list)

# ret_val : 1
# my_list : [2, 5, 13, 10]

heappush

heap모듈 heappush 함수는 힙 정렬된 상태를 유지하면서 데이터를 삽입시켜준다.

import heapq
my_list = [13, 2, 1, 5, 10]
heapq.heapify(my_list, -1)

# 가장 작은 원소인 -1이 가장 앞에 위치
print("남은 원소: ". ,my_list) # [-1, 2, 1, 5, 10, 13]

풀이

매번 리스트를 sort를 하고 데이터의 삽입 및 삭제가 빈번하게 일어나기 떄문에 heap구조를 활용합니다.

  1. 우선 들어오는 입력 리스트를 힙 정렬합니다.(heapsort)
  2. 최소값(m1)을 출력(heappop) 후 그 값이 K보다 모든 음식이 K보다 크기 떄문에 Stop.
  3. 만약 출력 후 입력 리스트의 길이가 0이라면, 모든 음식을 스코빌 지수 K이상으로 만들 수 없으므로 1을 리턴
  4. 다시 최소값(m2)을 출력하여 m1 + 2*m2 를 heap에 넣는다.(heappush) 이때 answer 에 1을 더해준다.
  5. 입력 리스트가 남아있을 때까지 2~4를 반복
  6. 최종 answer 를 출력한다
import heapq

def solution(scoville, K):
    answer = 0
    heapq.heapify(scoville)
    size = len(scoville)

    while True:
        m1 = heapq.heappop(scoville)
        if m1 >= K:
            return answer
        if len(scoville) == 0:
            return -1
        m2 = heapq.heappop(scoville)
        heapq.heappush(scoville, m1 + 2 * m2)
        answer += 1

    return answer