[파이썬] 프로그래머스 - 더 맵게

---

📌  Heap

Heap은 데이터를 정렬된 상태로 저장하기 위해서 사용하는 자료구조이다.

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

min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, min heap에서 가장 작은 값은 언제나 인덱스 0, 즉, 이진트리의 루트에 위치한다. 내부적으로 min heap 내의 모든 원소(k)는 항상 자식 원소들(2k+1, 2k+2) 보다 크기가 작거나 같도록 원소가 추가되고 삭제된다.

 

 

 

힙에 원소 추가

heapq 모듈의 heappush() 함수를 이용하여 힙에 원소를 추가할 수 있다. 첫 번째 인자는 원소를 추가할 대상 리스트이며 두 번째 인자는 추가할 원소를 넘긴다. (heapq 모듈에은 파이썬의 보통 리스트를 마치 최소 힙처럼 다룰 수 있도록 도와준다.)

 

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


[1, 3, 7, 4]

 

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

 

 

힙에서 원소 삭제

heapq 모듈의 heappop() 함수를 이용하여 힙에서 원소를 삭제할 수 있다. 원소를 삭제할 대상 리스트를 인자로 넘기면, 가장 작은 원소를 삭제 후에 그 값을 리턴하게 된다.

 

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


1
[3, 4, 7]

...

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


3
[4, 7]

 

최대 힙

파이썬의 heapq 모듈은 최소 힙(min heap)을 기능만을 동작하기 때문에 최대 힙(max heap)으로 활용하려면 약간의 요령이 필요하다. 힙에 튜플(tuple)을 원소로 추가하거나 삭제하면, 튜플 내에서 맨 앞에 있는 값을 기준으로 최소 힙이 구성되는 원리를 이용한다.

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

 

 

 

---

 

 

 

 

📌  풀이 

heap의 0번 인덱스가 가장 작은 원소이므로 이 원소가 K보다 작아질 때까지 while문을 수행한다.

while문 안에서는 가장 heappop을 통해 가장 작은 2개 원소를 꺼내 주어진 로직대로 처리하고 heappush를 통해 넣어준다. 그 값은 heap의 원칙에 따라 자동으로 크기에 맞춰 heap에 들어가게 된다.

 

또한 예외처리를 위해

try except을 통해 모든 음식을 K이상으로 만들 수 없을 때 answer에 -1을 넘겨줘 리턴한다.

 

 

코드

import heapq

def solution(scoville, K):
    scoville.sort()
    heapq.heapify(scoville)
    answer = 0
    try:
        while scoville[0] < K:
            a=heapq.heappop(scoville)
            b=heapq.heappop(scoville)
            heapq.heappush(scoville, a + b * 2)
            answer += 1
    except:
        answer = -1

    return answer