[프로그래머스] level2 | 더 맵게 - 파이썬(Python) | 힙(Heap)

https://school.programmers.co.kr/learn/courses/30/lessons/42626

프로그래머스

---

처음에는 리스트를 sorting 해 준 뒤에 while문을 동작 시키려고 했으나 효율성 테스트에서 실패가 떴다.

어떻게 시간복잡도를 개선할 수 있는지 찾아보니 heapq라는 아주 좋은 도구가 있었다.

 

heapq는 이미 힙구조라고 가정한 리스트를 조작하는 함수 모음으로 일반 리스트를 힙 구조로 사용하기 위해서는 heapq를 import 해 준뒤 heapq.heapify(scoville)을 해 줘야만 이 리스트를 이제부터 최소 힙으로 쓰겠다고 정의하게 된다.

import heapq

arr = [5, 3, 8, 1, 9]
heapq.heapify(arr)
print(arr)

위 코드를 실행하면 [1, 3, 8, 5, 9] 이렇게 출력이 된다.

이유는 heapq가 기대하는 최소 힙 구조는 모든 부모 노드는 자식 노드보다 작거나 같다는 것이다. 그래서 항상 heap[0[이 가장 작은 값으로 보장이 된다.

        1
      /   \
     3     8
    / \
   5   9

하지만 형제끼리는 크기 순서가 보장되지 않는다.

 

 

또 heapq.heapify(arr)는 리시트 전체를 O(n) 시간에 최소 힙 구조로 재배열한다.

처음 내가 시도했던 것 처럼 매번 정렬을 해주게 되면 O(n log n) 또는 O(n²) 시간 복잡도가 된다. 때문에 heapq를 사용해야만 효율성 테스트에서 통과할 수 있게 된다.

 

import heapq

heapq.heapify(arr)      # 1힙으로 변환
heapq.heappop(arr)      # 최소값 사용
heapq.heappush(arr, x)  # 다시 삽입

위 패턴은 이중 우선순위 큐 등에서도 거의 그대로 재사용된다.

 

이를 이용해 최종적으로 문제를 풀게 되면 코드가 아래와 같다.

import heapq

def solution(scoville, K):
    heapq.heapify(scoville)
    
    answer = 0
    mix = 0
    
    while scoville[0] < K:
        if len(scoville) < 2:
            return -1
        mix = heapq.heappop(scoville) + (heapq.heappop(scoville)*2)
        heapq.heappush(scoville, mix)
        
        answer += 1    
        
    return answer