[Level 2] 더 맵게

풀이

import heapq


def solution(scoville, k):
    heapq.heapify(scoville)  # scoville 리스트를 힙 구조로 만들어준다

    cnt = 0  # while 문 반복 횟수
    while scoville[0] < k:  # heap의 최소값이 k보다 커지면 break
        try:
            heapq.heappush(
                scoville, heapq.heappop(scoville) + (heapq.heappop(scoville) * 2)
            )
        except IndexError:  # 모든 음식의 스코빌 지수를 K 이상으로 만들 수 없는 경우
            return -1
        cnt += 1

    return cnt

인자로 들어온 scoville 리스트를 heapq.heapify()로 힙 구조로 만들어준다.

공식에 맞게 스코빌 지수를 구해주고, scoville에 원소가 남아 있지 않게 된다면 K 스코빌을 만들 수 없다는 의미이므로 -1을 리턴해준다.

Python heapq 모듈

힙 자료구조

heapq 모듈은 이진 트리(binary tree) 기반의 최소 힙(min heap) 자료구조를 제공한다.

min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, 작은 값은 언제나 인덱스 0, 다시 말해 이진 트리의 루트에 위치한다.

모듈 임포트

import heapq

내장 모듈이기 떄문에 위와 같이 import 해주면 된다.

최소 힙 생성

heap = []

heapq 모듈에서는 파이썬의 리스트를 마치 최소힙처럼 다룰 수 있게 해준다.

위와 같이 빈 리스트를 생성해 놓은 다음 heapq 모듈의 함수를 호출할 때마다 이 리스트에 인자를 넘겨야 한다. 다시 말해 파이썬에서는 heapq 모듈을 통해서 원솔르 추가하거나 삭제한 리스트가 그냥 최소 힙이다.

힙에 원소 추가

heapq 모듈의 heappush() 함수를 이용하여 힙에 원소를 추가할 수 있다.

heapq.heappush(heap, 4)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3)
print(heap)

[1, 3, 7, 4]

가장 작은 원소인 1은 인덱스 0에 위치하게 되며, 인덱스 1(=k)에 위치한 3은 인덱스 3(=2k+1)에 위치한 4보다 크므로 힙의 공식을 만족한다.

내부적으로 이진 트리에 원소를 추가하는 heappush() 함수는 $O(logN)$의 시간 복잡도를 가진다.

힙에서 원소 삭제

heapq 모듈의 heappop() 함수를 이용하여 힙에서 원소를 삭제 할 수 있다.

리스트를 heappop의 인자로 넣으면, 가장 작은 원소 값을 삭제 후 그 값을 리턴해준다.

print(heapq.heappop(heap))
print(heap)

1
[3, 4, 7]

print(heapq.heappop(heap))
print(heap)

3
[4, 7]

내부적으로 이진 트리로부터 원소를 삭제하는 heappop() 함수도 $ㅒO(logN)의 시간 복잡도를 가진다.

최소값 삭제하지 않고 얻기

힙에서 값을 삭제하지 않고 단순히 값을 얻어오기만 하려면 인덱스를 사용하면 된다.

print(heap[0])

여기서 주의할 점은 인덱스 0에 가장 작은 우너소가 있다고 해서, 인덱스 1에 두 번째로 작은 원소, 인덱스 2에 세 번쨰로 작은 원소가 있다는 보장이 없다는 것이다.

기존 리스트를 힙으로 변환

이미 원소가 들어있는 리스트를 힙으로 변환하려면 heapq 모듈의 heapify() 함수를 사용하면 된다.

heap = [4, 1, 7, 3, 8, 5]
heapq.heapify(heap)
print(heap)

[1, 3, 5, 4, 8, 7]

heapify() 함수에 리스트를 인자로 넘기면 리스트 내부의 원소들이 힙 구조에 맞게 재배치 되며 최소값은 0번째 인덱스에 위치하게 된다.

즉 비어 있는 리스트를 생성한 후 heappush() 함수로 원소를 한씩 추가한 작업과 같기 떄문에 heapify() 함수의 성능은 리스트의 원소 수와 비례하며 $O(N)$의 시간 복잡도를 가진다.

최대 힙

heapq 모듈은 기본적으로 최소 힙(min heap)의 기능만 제공하므로 최대 힙(max heap)을 만들려면 약갼의 요령이 필요하다.

힙에 원소를 추가할 때 2개의 값을 가지는 튜플(tuple)를 주면 앞의 원소를 기준으로 최소 힙이 구성되는 원리를 이용하는 것이다.

따라서 최대 힙을 만들려면 각 값에 대한 우선순위를 넣어주고, (우선 순위, 값) 구조로 힙에 추가하거나 삭제를 해주면 된다.
그리고 힙에서 값을 읽어올 떄는 각 튜플에서 인덱스 1에 있는 값을 가져오면 된다.

import heapq

nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
  heapq.heappush(heap, (-num, num))  # (우선 순위, 값)

while heap:
  print(heapq.heappop(heap)[1])  # index 1

8
7
5
4
3
1

K번째 최소값/최대값

import heapq

def kth_smallest(nums, k):
  heap = []
  for num in nums:
    heapq.heappush(heap, num)

  kth_min = None
  for _ in range(k):
    kth_min = heapq.heappop(heap)
  return kth_min

print(kth_smallest([4, 1, 7, 3, 8, 5], 3))

힙 정렬

import heapq

def heap_sort(nums):
  heap = []
  for num in nums:
    heapq.heappush(heap, num)

  sorted_nums = []
  while heap:
    sorted_nums.append(heapq.heappop(heap))
  return sorted_nums

print(heap_sort([4, 1, 7, 3, 8, 5]))

힙에 원소를 넣어준 후 하나씩 빼면 정렬이 된다.

자료 출처

• [파이썬] heapq 모듈 사용법