파이썬의 heapq 모듈

알고리즘 문제를 풀다 보면, heapq 모듈을 자주 만나게 된다.

특히 우선순위가 높은 값(또는 가장 낮은 값)을 빠르게 꺼내야 하는 상황에서 매우 유용하다.

 

1. heapq란?

heapq는 파이썬에서 힙(heap) 자료구조를 다룰 수 있게 해주는 내장 모듈이다.

힙은 완전 이진 트리 기반의 자료구조로,
가장 중요한 특징은 최솟값 또는 최댓값을 빠르게 꺼낼 수 있다는 점이다.

파이썬의 heapq는 기본적으로 최소 힙(min heap) 으로 동작한다.

 

즉, 가장 작은 값이 항상 맨 앞에 오도록 관리된다.

예를 들어 이런 상황에서 유용하다.

• 가장 작은 수를 반복해서 꺼내야 할 때
• 작업 우선순위가 높은 것부터 처리해야 할 때
• 다익스트라 알고리즘처럼 최소 비용을 계속 선택해야 할 때
• 실시간으로 값이 추가되면서 최솟값을 빠르게 확인해야 할 때

2. 힙(Heap)이란?

힙은 단순히 “정렬된 자료구조”라고 생각하면 조금 다르다.

많이 헷갈리는 부분인데, 힙은 전체가 완전히 정렬되어 있는 구조가 아니다.
대신 부모와 자식 사이의 대소 관계만 보장한다.

 

최소 힙의 경우:

• 부모 노드의 값이 자식 노드의 값보다 작거나 같다.
• 따라서 루트에는 항상 가장 작은 값이 온다.

예를 들어 아래와 같은 값들이 최소 힙으로 들어있다면:

[1, 3, 2, 7, 6, 5]

이 리스트는 전체가 정렬된 상태는 아닐 수 있다.
하지만 맨 앞의 원소 1은 항상 최소값이다.

 

즉, 힙은

• 최솟값 찾기는 빠르고
• 전체 정렬 상태 유지는 하지 않는

중간 성격의 자료구조라고 보면 된다.

 

3. 리스트 정렬 대신 힙을 쓰는 이유

예를 들어 리스트에 숫자를 넣고 매번 정렬해서 가장 작은 값을 꺼낸다고 해보자.

arr = [5, 2, 8, 1]
arr.sort()
print(arr.pop(0))

이 방식도 가능은 하지만 비효율적이다.

 

왜냐하면:

• 새 값이 들어올 때마다 다시 정렬하면 비용이 크고 ( O(logN) )
• pop(0)은 앞 원소를 꺼낸 뒤 나머지를 전부 당겨야 해서 느리기 때문이다 ( O(n) )

반면 힙은:

• 삽입: O(log n)
• 최소값 삭제: O(log n)
• 최소값 조회: O(1)

으로 동작하므로,
값이 자주 추가되고 최솟값을 자주 꺼내는 문제에서 훨씬 적합하다.

 

4. heapq 기본 사용법

import heapq

heap = []

heapq.heappush(heap, 4)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)
heapq.heappush(heap, 3)

print(heap)         # 내부 힙 구조
print(heapq.heappop(heap))  # 1
print(heapq.heappop(heap))  # 3
print(heapq.heappop(heap))  # 4
print(heapq.heappop(heap))  # 7

출력 순서를 보면 항상 작은 값부터 나온다.

 

인덱스 관계는 다음과 같다.

• 부모 인덱스: (i - 1) // 2
• 왼쪽 자식 인덱스: 2 * i + 1
• 오른쪽 자식 인덱스: 2 * i + 2

즉, 트리를 실제로 노드 객체로 만들지 않고
배열 하나로 트리 구조를 표현하는 방식이다.

 

5. heapify()

• 리스트를 한 번에 힙으로 만들기

이미 값이 들어 있는 리스트가 있다면,
하나씩 heappush()를 하지 않고 heapify()로 한 번에 힙 구조로 바꿀 수 있다.

import heapq

arr = [4, 10, 3, 5, 1]
heapq.heapify(arr)

print(arr)
print(heapq.heappop(arr))   # 1

 

heapify()는 리스트 자체를 바로 힙 구조로 바꾼다.

 

6. 최대 힙 구현 방법

예를 들어 가장 큰 값을 먼저 꺼내고 싶다면, 값을 음수로 넣으면 된다.

import heapq

heap = []

heapq.heappush(heap, -4)
heapq.heappush(heap, -1)
heapq.heappush(heap, -7)
heapq.heappush(heap, -3)

print(-heapq.heappop(heap))  # 7
print(-heapq.heappop(heap))  # 4
print(-heapq.heappop(heap))  # 3
print(-heapq.heappop(heap))  # 1

 

7. 튜플과 함께 사용해서 우선순위 큐 만들기

heapq는 숫자만 넣는 것이 아니라 튜플도 넣을 수 있다.
이때 파이썬은 튜플을 앞 원소부터 차례대로 비교한다.

import heapq

heap = []

heapq.heappush(heap, (2, "청소"))
heapq.heappush(heap, (1, "공부"))
heapq.heappush(heap, (3, "게임"))

print(heapq.heappop(heap))  # (1, '공부')
print(heapq.heappop(heap))  # (2, '청소')
print(heapq.heappop(heap))  # (3, '게임')

 

같은 우선순위일 경우 첫 번째 원소가 같으면 두 번째 원소를 비교

하지만 만약 서로 비교 불가능한 타입이면 에러 발생 가능

 

8. 자주 쓰는 추가 함수들

• heapq.heappushpop(heap, item)

원소를 넣고 바로 하나를 꺼낸다.

import heapq

heap = [2, 3, 5]
heapq.heapify(heap)

result = heapq.heappushpop(heap, 1)
print(result)  # 1
print(heap)    # [2, 3, 5]

이 함수는:

1. 먼저 넣고
2. 그 다음 최소값을 꺼낸다

그래서 작은 값을 넣으면 바로 빠질 수도 있다.

• heapq.heapreplace(heap, item)

최소값을 먼저 꺼내고, 새 값을 넣는다.

import heapq

heap = [2, 3, 5]
heapq.heapify(heap)

result = heapq.heapreplace(heap, 1)
print(result)  # 2
print(heap)    # [1, 3, 5]

• heapq.nlargest(n, iterable)
• heapq.nsmallest(n, iterable)

가장 큰 값 n개, 혹은 가장 작은 값 n개를 구할 수 있다.

import heapq

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

print(heapq.nsmallest(2, arr))  # [1, 3]
print(heapq.nlargest(2, arr))   # [9, 7]

 

9. 시간 복잡도 정리

heappush() : O(logN)

heappop() : O(logN)

최소값 확인 : heap[0] : O(1)

heapify() : O(N)