[Python] 파이썬의 heapq 모듈: 힙(Heap) 자료구조 활용 : 우선순위 큐, 최대 최소 등

 

우리는 종종 데이터를 정렬하거나 최소/최대값을 빠르게 찾아야 할 때가 있습니다. 파이썬의 heapq 모듈은 이러한 작업을 위해 사용되는 유용한 도구 중 하나입니다. 이번 포스팅에서는 heapq가 무엇인지, 어떤 함수들이 있는지, 그리고 실제 예제를 통해 어떻게 사용하는지 살펴보겠습니다.

힙(Heap)이란?

힙은 특별한 종류의 이진 트리로, 부모 노드가 자식 노드보다 작거나 큰 값을 가지는 자료구조입니다. 보통은 최소 힙(min heap)이나 최대 힙(max heap)으로 사용됩니다. 최소 힙은 부모 노드가 항상 자식 노드보다 작거나 같은 값을 가지며, 최대 힙은 부모 노드가 항상 자식 노드보다 크거나 같은 값을 가집니다.

heapq란?

heapq 모듈은 이진트리 기반의 최소 힙 자료구조를 제공하는 파이썬 표준 라이브러리입니다.

효율적인 우선순위 큐(priority queue) 구현할 수 있습니다. 이 모듈을 사용하면 원소들의 집합을 우선순위 큐 자료구조로 다룰 수 있습니다.

heapq가 PriorityQueue보다 실행시간이 적게 걸려 우선순위 큐 구현할 때 많이 사용합니다.

사용법

heapq 모듈은 파이썬의 리스트를 최소 힙(min-heap)으로 다룹니다. 즉, 리스트의 첫 번째 원소는 항상 최솟값입니다. heapq 모듈의 함수를 사용하여 리스트를 힙으로 변환하거나, 원소를 추가하고 삭제할 수 있습니다.

heapq 모듈 주요 함수

1. heapify(iterable): 주어진 리스트를 힙으로 변환합니다.
2. heappush(heap, item): 힙에 요소를 추가합니다.
3. heappop(heap): 힙에서 최소값을 제거하고 반환합니다.
4. heappushpop(heap, item): 요소를 힙에 추가한 후 최소값을 반환합니다.
5. heapreplace(heap, item): 최소값을 제거하고 새 요소를 추가합니다.

예제로 알아보는 heapq 사용법

from heapq import heappush, heappop, heapify, heappushpop,  heapreplace

# 힙 생성
heap = []

# 힙에 원소 추가 및 삭제
heappush(heap, 5)
heappush(heap, 3)
heappush(heap, 7)
heappush(heap, 1)
print(heap) # [1, 3, 7, 5]

print(heappop(heap)) # 1
print(heap)          # [3, 5, 7]

# 최소값 삭제하지 않고 얻기
print(heap[0]) # 3
print(heap)    # [3, 5, 7]

# 기존 리스트 힙으로 변환
heap2 = [3, 5, 1, 5, 2, 6]
heapify(heap2) 
print(heap2) # [1, 2, 3, 5, 5, 6]

# 요소 추가 후 최소값 반환
min_value = heappushpop(heap2, 7)
print(min_value) # 1
print(heap2)     # [2, 5, 3, 5, 7, 6]
# 아래처럼 넣는 것까지 포함하여 최소값을 구하고 반환한다
min_value = heappushpop(heap2, 1)
print(min_value) # 1
print(heap2)     # [2, 5, 3, 5, 7, 6]

# 최소값 제거 후 요소 추가
min_value = heapreplace(heap2, 10)
print(min_value) # 2
print(heap2)     # [3, 5, 6, 5, 7, 10]
# 아래처럼 최소값을 제거하고서 값을 넣는다
min_value = heapreplace(heap2, 1)
print(min_value) # 3
print(heap2)     # [1, 5, 6, 5, 7, 10]

heapq로 구현 가능한 것들

: 최대 힙, 우선순위 큐, n번째 최대/최소 값

최대 힙(max heap) 구현 방법 2가지

1. 음수를 사용

import heapq

def max_heap(iterable):
    heap = []
    for value in iterable:
        heapq.heappush(heap, -value)  # 음수를 사용하여 최대 힙 구현
    return [-heapq.heappop(heap) for _ in range(len(heap))]

# 테스트
data = [3, 1, 4, 1, 5, 9, 2, 6, 5]
result = max_heap(data)
print(result)  # 출력: [9, 6, 5, 5, 4, 3, 2, 1, 1]

2. 우선순위를 사용

import heapq

def max_heap(iterable):
    heap = []
    for value in iterable:
        heapq.heappush(heap, (-value, value))  # (우선 순위, 값)
        # 이러면 우선순위에 따라 정렬되어 8, 3이 -8, -3으로 -8이 더 우선이라 
        # 인덱스[1]로만 봤을 때는 8이 더 앞에 오게 됨
    return [heapq.heappop(heap)[1] for _ in range(len(heap))]

# 테스트
data = [3, 1, 4, 1, 5, 9, 2, 6, 5]
result = max_heap(data)
print(result)  # 출력: [9, 6, 5, 5, 4, 3, 2, 1, 1]

 

우선순위 큐

우선순위가 숫자가 작을 때 더 우선이면 min heap(default) 사용하여 튜플로 (우선순위, 아이템) 넣기 - 숫자 큰 것이 더 우선이면 max heap으로 동일 구현

import heapq

# 우선순위 큐를 생성
priority_queue = []

# 아이템과 우선순위 추가
heapq.heappush(priority_queue, (5, 'task1'))
heapq.heappush(priority_queue, (1, 'task2'))
heapq.heappush(priority_queue, (3, 'task3'))

# 아이템 추출
print(heapq.heappop(priority_queue)[1])  # 출력: task2
print(heapq.heappop(priority_queue)[1])  # 출력: task3
print(heapq.heappop(priority_queue)[1])  # 출력: task1

 

n번째 최소값/최대값

1. n번째 최소값 : min heap(default)에서 n번 pop

import heapq

def nth_smallest(nums, n):
   # 배열을 힙으로 만들기 1
    heap = nums
    heapq.heapify(heap) 

   # 배열을 힙으로 만들기 2
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
    #######

    nth_min = None 
    for _ in range(n):
        nth_min = heapq.heappop(heap)
    return nth_min

print(nth_smallest([4, 1, 7, 3, 8, 5], 3)) # 4

2. n번째 최소값 : heapq 모듈에는 이미 이러한 용도의 nsmallest() 함수 있다!

import heapq

nth_small = heapq.nsmallest(3, [4, 1, 7, 3, 8, 5])[-1] 
print(nth_small) # 4

3. n번째 최대값: max heap 만들고 n번 pop

import heapq

def nth_maximum(nums, n):
    heap = []
    for num in nums:
        heapq.heappush(heap, -num)

    nth_max = None
    for _ in range(n):
        nth_max = heapq.heappop(heap)

    return -nth_max


print(nth_maximum([4, 1, 7, 3, 8, 5], 3)) # 5

4. n번째 최대값: heapq 모듈에는 이미 이러한 용도의 nlargest() 함수 있다!

import heapq

nth_max = heapq.nlargest(3, [4, 1, 7, 3, 8, 5])[-1]
print(nth_max) # 5

 

heapq 모듈의 nsmallest와 nlargest

💡 `heapq` 모듈의 `nsmallest`와 `nlargest` 함수는 주어진 iterable에서 가장 작은 또는 가장 큰 n개의 요소를 찾아 반환

nsmallest 함수

nsmallest(n, iterable, key=None): 주어진 iterable에서 가장 작은 n개의 요소를 찾아 반환

• n: 반환할 요소의 개수를 나타내는 정수 값
• iterable: 요소를 탐색할 iterable 객체
• key (선택적): 정렬 기준을 제공하는 함수. 기본값은 None이며, 이 경우 요소들은 기본적으로 작은 순서로 정렬.

 

nlargest 함수

nlargest(n, iterable, key=None): 주어진 iterable에서 가장 큰 n개의 요소를 찾아 반환

• n: 반환할 요소의 개수를 나타내는 정수 값
• iterable: 요소를 탐색할 iterable 객체
• key (선택적): 정렬 기준을 제공하는 함수. 기본값은 None이며, 이 경우 요소들은 기본적으로 큰 순서로 정렬.

예제 코드

import heapq

# nsmallest 예제
data = [3, 1, 4, 1, 5, 9, 2, 6, 5]
smallest = heapq.nsmallest(3, data)
print("가장 작은 3개의 요소:", smallest)  # 출력: [1, 1, 2]

# nlargest 예제
largest = heapq.nlargest(3, data)
print("가장 큰 3개의 요소:", largest)    # 출력: [9, 6, 5]

# 사용자 지정 key 예제
students = [
    {'name': 'Alice', 'grade': 90},
    {'name': 'Bob', 'grade': 85},
    {'name': 'Charlie', 'grade': 88},
    {'name': 'David', 'grade': 92},
]
top_students = heapq.nlargest(2, students, key=lambda x: x['grade'])
print("성적 상위 2명:", top_students)   
# 출력: [{'name': 'David', 'grade': 92}, {'name': 'Alice', 'grade': 90}]