Python에서 heapq로 최대힙, 최소힙 구현하기

Python의 내장 모듈인 heapq는 힙(Heap) 자료 구조를 제공하며, 기본적으로 최소힙(Min-Heap)으로 동작합니다. 최소힙은 가장 작은 요소를 항상 루트 노드로 유지하는 완전 이진 트리입니다.

하지만 최대힙(Max-Heap)은 지원하지 않으므로, 특정 트릭을 사용하여 최대힙을 구현해야 합니다.

1. 최소힙(Min-Heap) 구현

기본 동작

• heapq는 리스트를 힙처럼 다룰 수 있도록 제공합니다.
• heapq.heappush(heap, item): 힙에 새로운 요소를 추가합니다.
• heapq.heappop(heap): 힙에서 가장 작은 요소를 제거하고 반환합니다.

예제 코드

import heapq

# 최소힙 생성
min_heap = []

# 요소 추가
heapq.heappush(min_heap, 3)
heapq.heappush(min_heap, 1)
heapq.heappush(min_heap, 4)
heapq.heappush(min_heap, 2)

print("Min-Heap:", min_heap)  # [1, 2, 4, 3]

# 가장 작은 요소 제거 및 반환
smallest = heapq.heappop(min_heap)
print("Popped element:", smallest)  # 1
print("After pop:", min_heap)       # [2, 3, 4]

결과

1. 요소를 추가하면 내부적으로 힙의 속성을 유지합니다.
2. heappop을 호출하면 가장 작은 값을 제거하고 반환합니다.

2. 최대힙(Max-Heap) 구현

Python의 heapq는 최소힙만 지원하므로, 음수값을 사용하여 최대힙처럼 동작하도록 만들 수 있습니다:

• 삽입 시 값에 -1을 곱해 음수로 저장.
• 제거 시 다시 -1을 곱해 원래 값으로 복구.

예제 코드

import heapq

# 최대힙 생성
max_heap = []

# 요소 추가 (음수로 변환하여 추가)
heapq.heappush(max_heap, -3)
heapq.heappush(max_heap, -1)
heapq.heappush(max_heap, -4)
heapq.heappush(max_heap, -2)

# 힙 내부 상태 확인 (음수로 저장됨)
print("Max-Heap (inverted):", max_heap)  # [-4, -2, -3, -1]

# 가장 큰 요소 제거 및 반환 (다시 양수로 변환)
largest = -heapq.heappop(max_heap)
print("Popped element:", largest)  # 4
print("After pop:", [-x for x in max_heap])  # [3, 1, 2]

결과

1. 힙에 삽입할 때 음수로 변환하여 최대값을 루트로 유지합니다.
2. 제거 시 음수를 다시 양수로 변환해 원래 값을 반환합니다.

3. 최소힙과 최대힙을 동시에 사용하는 경우

어떤 문제에서는 최소힙과 최대힙을 함께 사용해야 할 때가 있습니다. 예를 들어, 중간값을 빠르게 구하기 위해 최소힙과 최대힙을 사용하는 경우가 있습니다.

예제: 중간값 유지

import heapq

min_heap = []  # 최소힙 (큰 절반 저장)
max_heap = []  # 최대힙 (작은 절반 저장, 음수로 변환)

def add_number(num):
    if not max_heap or num <= -max_heap[0]:
        heapq.heappush(max_heap, -num)  # 최대힙에 추가
    else:
        heapq.heappush(min_heap, num)  # 최소힙에 추가
    
    # 크기 조정 (최대힙 크기가 최소힙보다 크거나 같게 유지)
    if len(max_heap) > len(min_heap) + 1:
        heapq.heappush(min_heap, -heapq.heappop(max_heap))
    elif len(min_heap) > len(max_heap):
        heapq.heappush(max_heap, -heapq.heappop(min_heap))

def find_median():
    if len(max_heap) == len(min_heap):
        return (-max_heap[0] + min_heap[0]) / 2
    return -max_heap[0]

# 예제 사용
add_number(1)
add_number(5)
add_number(2)
add_number(10)
add_number(6)
print("Median:", find_median())  # Median: 5

4. 힙 정렬 구현

힙을 활용하면 정렬도 쉽게 구현할 수 있습니다.

예제 코드

import heapq

def heap_sort(iterable):
    heap = []
    for value in iterable:
        heapq.heappush(heap, value)  # 요소를 힙에 추가
    
    sorted_list = [heapq.heappop(heap) for _ in range(len(heap))]
    return sorted_list

nums = [5, 1, 8, 3, 7, 2]
print("Heap Sort:", heap_sort(nums))  # [1, 2, 3, 5, 7, 8]

5. 정리

기능                               최소힙                     최대힙

사용 방법	heappush, heappop	heappush(-value), -heappop()
루트 노드 값	가장 작은 값	가장 큰 값
활용 사례	우선순위 큐, K번째 작은 값 찾기	최대값 추적, K번째 큰 값 찾기
예제 코드	기본 제공	음수 변환을 통해 구현