[Python] 파이썬 힙 - heapq 라이브러리 nlargest / nsmallest

heapq 모듈은 파이썬에서 우선순위 큐 알고리즘이라고도 하는 힙(heap) 큐 알고리즘의 구현을 제공한다.

그 중 유용하게 사용할 수 있는 nlargest / nsmallest 함수를 알아보자!

 

Heapq.nsmallest()

heapq.nsmallest(n, iterable, key=None)

Heap의 아이템을 작은 순서대로 n개를 리스트 형태로 리턴한다.

(n 개의 가장 작은 요소로 구성된 리스트를 반환)

매개변수 n의 값을 힙의 길이로 전달할 때, 힙을 오름차순으로 정렬한 것과 같은 결과를 제공한다.

 

Heapq.nlargest()

heapq.nlargest(n, iterable, key=None)

Heap의 아이템을 큰 순서대로 n개를 리스트 형태로 리턴한다.

(n 개의 가장 큰 요소로 구성된 리스트를 반환)

힙을 내림차순으로 정렬한 것과 같은 결과를 제공한다.

 

예시

import heapq

heap = [-45, 33, 20, -90, 41]
heapq.heapify(heap)

# 1. 힙의 아이템을 작은 순대로 3개 출력
print('작은 수 3개 : ', heapq.nsmallest(3, min_h))

# 2. 힙의 아이템을 큰 순대로 3개 출력
print('큰 수 3개 : ', heapq.nlargest(3, heap))

최소 힙을 생성하고, 작은 수 3개와 큰 수 3개를 출력하는 예제이다.

 

출력 결과

---

활용 예시

import heapq

# 1. 최소힙 생성
min_h = [-45, 33, 20, -90, 41]
heapq.heapify(min_h)

# 2. 최대값 삭제
min_h = heapq.nlargest(len(min_h), min_h)[1:]
heapq.heapify(min_h)

수가 작은 순서대로(오름차순) 정렬되어있는 최소 힙에서 최대 값 1개를 삭제하려고 할 때, 위와 같은 방식으로 삭제할 수 있다.

1. 최소 힙을 nlargest 함수를 사용하여 내림차순 정렬한 리스트를 구한 다음,

2. 가장 첫번째 값(= 최대값)을 슬라이싱 한다.

3. 최대값을 제외한 리스트를 다시 최소 힙으로 만든다. (heapify)

 

결과

# 힙의 초기값
[-90, -45, 20, 33, 41]

# 최댓값 삭제 후
[-90, 20, -45, 33]​

최소 힙에서 최대값인 41이 삭제되는 결과를 확인할 수 있다.