파이썬 heapq 사용법

heapq는 Binary Heap 기반 우선순위 큐입니다.

구조는 다음과 같습니다.

heapq 모듈 소개

파이썬의 heapq 모듈은 리스트를 기반으로 힙을 구현합니다. 이 모듈을 사용하면 데이터를 효율적으로 관리하고, 가장 작은 원소를 쉽게 찾고 제거할 수 있습니다.

기본 사용법

import heapq

1. 리스트를 힙으로 변환하기

리스트를 힙으로 변환하려면 heapify 함수를 사용합니다.

data = [5, 7, 9, 1, 3]
heapq.heapify(data)
print(data)  # 결과: [1, 3, 9, 7, 5]

2. 원소 추가하기

리스트에 원소를 추가하고 싶다면 heappush 함수를 사용하세요.

heapq.heappush(data, 4)
print(data)  # 결과: [1, 3, 4, 7, 5, 9]

3. 최소 원소 제거 및 반환하기

가장 작은 원소를 제거하고 싶다면 heappop 함수를 사용합니다.

min_element = heapq.heappop(data)
print(min_element)  # 결과: 1
print(data)  # 결과: [3, 4, 9, 7, 5]

4. 여러 개의 원소 추가하기

원소를 추가하면서, 동시에 힙의 구조를 유지하고 싶다면 heappushpop을 사용해보세요.

min_element = heapq.heappushpop(data, 2)
print(min_element)  # 결과: 1
print(data)  # 결과: [2, 3, 9, 7, 5]

5. 최소 원소를 유지하면서 원소 추가하기

heapreplace 함수는 가장 작은 원소를 제거하고 새로운 원소를 추가합니다.

min_element = heapq.heapreplace(data, 6)
print(min_element)  # 결과: 2
print(data)  # 결과: [3, 4, 6, 7, 5]

6. n개의 최소 원소 찾기

nsmallest 함수를 사용하면 리스트에서 n개의 가장 작은 원소를 쉽게 찾을 수 있습니다.

smallest = heapq.nsmallest(3, data)
print(smallest)  # 결과: [3, 4, 5]

7. n개의 최대 원소 찾기

반대로, nlargest 함수를 사용하여 n개의 가장 큰 원소를 찾을 수도 있습니다.

largest = heapq.nlargest(3, data)
print(largest)  # 결과: [7, 6, 9]

예제 코드

아래는 heapq 모듈의 기본적인 사용 예제입니다.

import heapq

# 데이터 초기화
data = [10, 20, 15, 30, 40]

# 힙으로 변환
heapq.heapify(data)

# 원소 추가
heapq.heappush(data, 5)
print(data)  # 결과: [5, 20, 10, 30, 40, 15]

# 최소 원소 제거
print(heapq.heappop(data))  # 결과: 5
print(data)  # 결과: [10, 20, 15, 30, 40]