[Python] Heap과 heapq 모듈

heap 에 대한 설명은 https://brownbears.tistory.com/391에서 했지만 여기서는 더 자세하게 설명을 합니다.

Heap은 최댓값, 최솟값을 찾아내는 연산을 빠르게 하기 위해 고안되었으며 완전 이진 트리 (Complete Binary Tree)를 기본으로 가집니다. 즉, 데이터가 입력이 되면 가장 말단에서 왼쪽부터 차례대로 채워져 나갑니다.

최대 힙(max heap)은 부모의 노드가 자식 노드의 값과 같거나 더 크며 최소 힙(min heap)은 부모의 노드가 자식 노드의 값과 같거나 더 작습니다. 힙과 이진 탐색 트리 (binary search tree)이 쉽게 헷갈리는데 이진 탐색 트리의 경우, 부모 노드의 값보다 작으면 왼쪽, 크면 오른쪽 이라는 규칙이 있지만 힙의 경우엔 없습니다. 즉, 부모 노드를 기준으로 왼쪽 노드가 오른쪽 노드보다 작을 수도 있습니다.

heap을 사용하는 가장 큰 이유는 최댓값, 최솟값을 찾기 위함이며 데이터 삽입, 삭제는 모두 O(log₂N)의 시간 복잡도가 걸립니다.

여기서는 최소 힙을 기준으로 설명을 진행합니다. (최대 힙과 삽입, 삭제가 동일하기 때문에 중복 내용을 작성하지 않습니다)

Insert

데이터 삽입은 아래와 같은 순서를 지킵니다.

1. 데이터 삽입은 왼쪽에서 오른쪽 순서로 완전 이진 트리를 구성
2. 입력된 데이터와 부모 노드의 값을 비교
3. 부모보다 작다면(최소힙 기준) 서로의 자리를 변경
4. 2~3번을 부모보다 클때까지 반복

아래는 [14, 10, 6, 8, 12, 4] 의 데이터 순으로 삽입해 최소 힙을 구성하는 예시입니다.

위 예시에서 10을 삽입한 다음, 부모와 비교해 부모보다 작을 경우, 위치를 서로 변경해 줍니다.

delete

힙에서 삭제는 루트 노드를 지우는 것이 일반적입니다.

데이터 삭제는 아래와 같은 순서를 지킵니다.

1. 루트노드 삭제
2. 트리의 말단에서 가장 오른쪽에 있는 노드를 루트로 승격
3. 루트 노드가 된 데이터와 두 자식 노드의 값을 비교하여 가장 작은 값과 위치 변경

아래는 위에서 만든 최소 힙에서 루트 노드를 삭제한 예시입니다.

heapq 모듈

파이썬 heapq 모듈은 최소 힙을 지원하는 모듈로 직접 최소 힙을 구현하지 않아도 되는 장점이 있습니다.

부모 노드의 인덱스를 1이라 할 때, 자식 노드의 인덱스는 다음과 같습니다.

• 부모 노드 인덱스 = 자식 인덱스 // 2
• 왼쪽 자식 노드 인덱스 = 부모 노드 인덱스 * 2
• 오른쪽 자식 노드 인덱스 = 부모 노드 인덱스 * 2 + 1

heapq 모듈은 인자로 넘겨받은 리스트를 재구성합니다. (반환값 X)

데이터 추가

import heapq

heapq.heappush(result, 14)
print(result)
# [14]

heapq.heappush(result, 10)
print(result)
# [10, 14]

heapq.heappush(result, 6)
print(result)
# [6, 14, 10]

heapq.heappush(result, 8)
print(result)
# [6, 8, 10, 14]

heapq.heappush(result, 12)
print(result)
# [6, 8, 10, 14, 12]

heapq.heappush(result, 4)
print(result)
# [4, 8, 6, 14, 12, 10]

데이터 삭제

최소, 최대 힙에서 데이터 삭제는 루트 노드를 삭제하는 것이 일반적이라 설명한 것처럼 heapq에서는 루트노드만 삭제할 수 있습니다.

heapq.heappop(result)
print(result)
# [6, 8, 10, 14, 12]

리스트 변환

기존 리스트를 최소 힙으로 변환은 다음과 같이 진행합니다.

_list = [14, 10, 6, 8, 12, 4]
heapq.heapify(_list)
print(_list)
# [4, 8, 6, 10, 12, 14]

위 heappush와 리스트의 결과가 다른 이유는 heapq의 내부적인 로직으로 인해 리스트를 순차적으로 진행하지 않기 때문입니다. heap에서는 최소, 최대값이 가장 처음 나오면 되기 때문에 이러한 결과는 문제가 되지 않습니다. heapify()는 O(N) 의 시간복잡도를 가집니다.

최대 힙 표현하기

heapq는 최소 힙만 지원하기 때문에 최대 힙을 만드려면 입력 값들에 음수를 씌우면 됩니다.

reverse_sign = lambda x: x*-1
max_heap = list(map(reverse_sign, _list))
heapq.heapify(max_heap)
max_heap = list(map(reverse_sign, max_heap))

# [14, 12, 6, 8, 10, 4]

힙 정렬

sorted()라는 좋은 정렬 함수가 내장되어 있지만 꼭 힙 정렬만 구현해야 한다면 아래와 같이 간단하게 구현을 할 수 있습니다.

import heapq

_list = [14, 10, 6, 8, 12, 4]
heapq.heapify(_list)

method1 = []

while _list:
    method1.append(heapq.heappop(_list))
print(method1)
# [4, 6, 8, 10, 12, 14]

또는 heapq의 함수를 이용해 사용할 순 있지만 추천을 하지 않습니다. heapq 함수 내부적으로도 구하고자 하는 리스트의 숫자와 입력된 리스트의 길이가 동일하면 sorted() 내장 함수를 쓰기 때문입니다.

import heapq

_list = [14, 10, 6, 8, 12, 4]
# heapq.heapify(_list)

method2 = heapq.nsmallest(len(_list), _list)
print(method2)
# [4, 6, 8, 10, 12, 14]