[Python] heapq(우선순위 큐) 사용법

파이썬의 heapq 라이브러리를 통해 손 쉽게 최소힙과 최대힙을 구현할 수 있다.

 

우선, heapq는 기본적으로 최소힙으로 구현되어있다.

 

즉, heapq의 heappush를 통해 값들을 삽입하면 해당 값들은 숫자가 가장 작은 순서대로 트리 구조로 값이 저장된다.

 

heapq의 연산을 사용하기 위해선 각 연산의 파라미터로 큐로 사용할 리스트와 원소를 넘겨주면 된다.

 

1. heappush - 값 추가 

import heapq

heap_q = []

heapq.heappush(heap_q, 3)
heapq.heappush(heap_q, 10)
heapq.heappush(heap_q, 1)
heapq.heappush(heap_q, 0)
heapq.heappush(heap_q, 4)

print(heap_q)

# [0, 1, 3, 10, 4]

이를 트리로 그려보면 다음과 같다.

 

우선순위 큐는 일반적으로 힙을 통해서 구현하기 때문에, 완전이진트리의 성질을 띠고 이를 이용해 특정 노드의 왼쪽 자식과 오른쪽 자식의 인덱스를 특정할 수 있다.

 

2. heappop - 값 삭제

import heapq

heap_q = []

heapq.heappush(heap_q, 3)
heapq.heappush(heap_q, 10)
heapq.heappush(heap_q, 1)
heapq.heappush(heap_q, 0)
heapq.heappush(heap_q, 4)

print(heapq.heappop(heap_q))

# 0

heappop 연산을 통해 우선순위 큐의 가장 우선순위가 높은 값(최소값)을 제거할 수 있다.

 

물론, 최소값이 제거되더라도 자동으로 힙 성질을 유지해준다.

 

import heapq

heap_q = []

heapq.heappush(heap_q, 3)
heapq.heappush(heap_q, 10)
heapq.heappush(heap_q, 1)
heapq.heappush(heap_q, 0)
heapq.heappush(heap_q, 4)

# 최소값 제거
heapq.heappop(heap_q)
print(heap_q)

# [1, 4, 3, 10]

 

3. heapify - 리스트를 힙으로 변환

heapify 연산을 통해서 아무런 정렬이 되지 않은 리스트의 원소들을 힙으로 변환할 수 있다.

import heapq

heap_q = [9, 20, 1, 8, 5, 0]

heapq.heapify(heap_q)
print(heap_q)

# [0, 5, 1, 8, 20, 9]

 

4. nlargest, nsmallest - 힙의 n개의 가장 큰 리스트, n개의 가장 작은 리스트

우선순위 큐에 저장된 원소들 중 n개의 가장 큰 리스트와 n개의 가장 작은 리스트를 반환한다.

import heapq

heap_q = []

heapq.heappush(heap_q, 3)
heapq.heappush(heap_q, 10)
heapq.heappush(heap_q, 1)
heapq.heappush(heap_q, 0)
heapq.heappush(heap_q, 4)

print(heap_q)

# heap_q에서 가장 큰 3개의 원소가 담긴 리스트
print(heapq.nlargest(n=3, iterable=heap_q))

# heap_q에서 가장 작은 3개의 원소가 담긴 리스트
print(heapq.nsmallest(n=3, iterable=heap_q))

# heap_q
# [0, 1, 3, 10, 4]

# heap_q의 nlargest
# [10, 4, 3]

# heap_q의 nsmallest
# [0, 1, 3]

 

최소힙을 최대힙으로 사용하기

처음 언급했듯, 파이썬의 heapq 라이브러리는 최소힙으로 구현되어있다.

 

하지만, 만약 이를 최대힙 즉, 가장 큰 수에 제일 높은 우선순위를 부여하고 싶다면 간단히 해당 수에 -1을 곱해 음수로 만들어주면 된다.

 

이후, pop할 때 다시 -1를 곱해 원래대로 양수로 만들어주면 최소힙을 최대힙으로 만들어 사용할 수 있다.

 

import heapq

heap_q = []

heapq.heappush(heap_q, -3)
heapq.heappush(heap_q, -10)
heapq.heappush(heap_q, -1)
heapq.heappush(heap_q, -0)
heapq.heappush(heap_q, -4)

print(-1 * heapq.heappop(heap_q))

# 10