[Python] heapq

데이터를 정렬된 상태로 저장하기 위해 사용하는 파이썬의 내장 모듈인 heapq에 대해 알아보자.

 

heapq 모듈은 binary tree 기반의 min heap 자료구조를 제공한다. min heap을 사용하면 원소들이 항상 정렬된 상태로 추가되고 삭제되며, min heap에서 가장 작은 원소의 값은 언제나 binary tree의 루트(인덱스 0)에 위치한다.

 

내부적으로 min heap 내의 모든 원소(k)는 항상 자식 원소들(2k+1, 2k+2) 보다 크기가 작거나 같도록 원소가 추가, 삭제된다.

 

아래의 그림으로 위의 조건을 만족시키는 간단한 min heap 구조를 보자.

 

우선 heapq를 사용하기 위해서는 다음과 같이 import 후 관련 함수를 사용할 수 있다.

import heapq

파이썬에서는 보통의 리스트를  min heap처럼 다룰 수 있도록 도와준다. 빈 리스트를 생성한 다음 heapq 모듈의 함수를 통해 추가하거나 삭제하면 해당 리스트가 min heap이 되는 것이다.

 

아래의 예제를 보자

min_heap = []

# 원소 추가
heapq.heappush(min_heap, 4)
heapq.heappush(min_heap, 1)
heapq.heappush(min_heap, 7)
heapq.heappush(min_heap, 2)

print(min_heap)

# [1, 2, 7, 4]

# 원소 삭제
print(heapq.heappop(min_heap))
print(min_heap)

# 1
# [2, 4, 7]

print(min_heap[0])

# 2

위와 같이 원소 추가 메서드로 heappush()를 사용하고, 원소 삭제 메서드로는 heappop()을 사용한다. heappush()를 사용한 경우 1번째(k) 인덱스의 2를 보면 2k+1(3번째) 인덱스의 원소인 4가 2보다 크거나 같기 때문에 heap의 공식을 만족하는 것을 볼 수 있고, heappop()을 사용하면 가장 작은 원소를 삭제 후에 그 값을 리턴한다.

 

heap에서 최솟값을 삭제하지 않고 단순히 읽기만 하려면 일반적으로 리스트의 원소에 접근하듯이 인덱스를 통해 접근하면 된다.

 

기존에 원소를 가지고 있는 리스트를 heap으로 변환하려면 heapq 모듈의 heapify() 메서드를 사용하면 된다.

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

print(list)

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

heapify를 통하면 기존의 리스트 원소들을 heap의 구조에 맞게 재배치한다.

 

 

참고 : https://www.daleseo.com/python-heapq/