[Python] 파이썬 heap(heapq 모듈)

 

프로그래머스에서 코딩 문제를 풀다가 heap을 써서 풀어야 만족할 수 있는 문제를 접했다...(반드시 모듈을 import해서 풀도록 유도하는 문제는 그리 좋은 것 같지 않지만....) 공부도 할겸 파이썬에서 heap을 사용하는 법을 정리해본다. 사실 정확히는 heapq이다. 풀이를 찾아보니 아래와 같이 코드를 구현하면 된다. 여기서 핵심인 heapify, heappop, heappush에 대해 알아보자.

import heapq as hq

def solution(scoville, K):

    hq.heapify(scoville)
    answer = 0
    while True:
        first = hq.heappop(scoville)
        if first >= K:
            break
        if len(scoville) == 0:
            return -1
        second = hq.heappop(scoville)
        hq.heappush(scoville, first + second*2)
        answer += 1  

    return answer

 

heapify() - heap 생성

---

역할은 주어진 리스트를 heap으로 만들어주는 것이다. heapify(리스트이름)과 같이 사용한다. 기본적으로 minheap을 생성한다.

 

heappop() - heap 원소 삭제

---

heap에서 pop을 하는데 minheap이니 deletemin()과 같다고 생각하면 된다. 즉, 힙에서 최솟값인 루트노드를 pop한다. heappop(리스트이름)과 같이 사용한다.

 

heappush() - heap 원소 추가

---

heap에 원소를 추가하는 것이다. 내부적으로는 heap에 원소를 추가하고 다시 heapness를 갖도록 재배치하는 과정을 거친다. heap에 원소를 추가할 때는 last node로 추가하고 부모노드와 비교하며 맞는 자리를 찾아가는 과정을 거치니 일반적으로 O(logn)의 시간복잡도를 갖는다. heappush(리스트이름, 추가할 값)과 같이 사용한다.

 

Maxheap 구현

---

좀 더 나아가서 maxheap(최대힙)을 구현하려면 어떻게 해야 할까? 아쉽게도 파이썬의 heapq가 기본적으로 제공하는 기능은 minheap이기 때문에 약간의 트릭을 써서 구현해야 한다. priority의 부호를 바꾸어 정렬하는 것이다. 예시는 아래와 같다.

heap_items = [1,3,5,7,9]

max_heap = []
for item in heap_items:
  heapq.heappush(max_heap, (-item, item))

print(max_heap)		# [(-9,9),(-7,7),(-5,5),(-3,3),(-1,1)]

 

참고자료

• littlefoxdiary.tistory.com/3