Python Heapq 모듈 사용법 (최소 힙, 최대 힙)

Rachel·2024년 5월 14일

heap

Algorithm

목록 보기

7/10

Heapq

Heap은 우선순위 큐(Priority Queues)를 구현하는데 널리 사용되는 자료 구조 중 하나.
파이썬에서는 heapq 모듈을 사용하면 된다.

우선순위가 높은 원소부터 먼저 추출된다.
-> 작업 스케줄링, 네트워크 패킷 라우팅, 이벤트 처리, 작업 예약, 우선순위에 따른 데이터 처리 등에 유용하게 활용됨.

사용법

heap 자료구조는 삽입, 반환시 항상 정렬을 유지한다.
heapq 모듈은 이진트리 기반의 최소 힙 자료구조로 최대 힙이나 최소 힙을 구현할 수 있다.

heap 만들기

import heapq

# 배열 초기화
li = [3, 7, 8, 9]

# list를 heap으로 변환
heap.heapify(li)

원소 넣기, 빼기

heappush(heap, ele): heap에 ele 넣기
heappop(heap): heap에서 가장 작은 원소 빼서 반환

import heapq

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

heapq.heapify(li)

# li에 4 넣기
heapq.heappush(li, 4)

# 1 반환
heapq.heappop(li)

동시에 원소 넣고 빼기

heappushpop(heap, ele): 동시에 ele 넣고, heap에서 가장 작은 값 빼서 반환
heapreplace(heap, ele): heap에서 가장 작은 값이 먼저 빠지고 반환된 후, ele가 push됨.

import heapq

li1 = [5, 1, 9, 4, 3]

li2 = [5, 7, 9, 4, 3]

heapq.heapify(li1)
heapq.heapify(li2)

# 2 삽입, 1 반환
heapq.heappushpop(li1, 2)

# 3 반환, 2 삽입
heapq.heapreplace(li2, 2)

가장 작은 값, 가장 큰 값 찾기

nlargest(k, iterable, key=fun): iterable에서 key를 만족하면서 가장 큰 K개 원소들을 반환
nsmallest(km iterable, key=fun): iterable에서 key를 만족하면서 가장 작은 k개 원소들을 반환

원소 삭제 x

import heapq

# initializing list
li1 = [6, 7, 9, 4, 3, 5, 8, 10, 1]

heapq.heapify(li1)

# 가장 큰 원소 3개 10, 9, 8 반환
heapq.nlargest(3, li1)

# 가장 작은 원소 3개 1, 3, 4 반환
heapq.nsmallest(3, li1)

가장 작은 원소 1개 찾고 싶으면 heappop()을 이용해도 된다.

heapq.heappop(li)

Heap 활용 - 알고리즘 문제

힙을 배웠다면 활용해서 알고리즘 문제를 풀어보자!

백준 1927번 최소 힙

풀이

import sys
import heapq
input = sys.stdin.readline

n = int(input())
que = []

for _ in range(n):
    x = int(input())
    if x == 0:
        if len(que) == 0:
              print(0)
        else:
            print(heapq.heappop(que))
    else:
       heapq.heappush(que, x)

heappop으로 쉽게 가장 작은 원소를 []에서 삭제, 반환할 수 있다.

백준 11279번 최대 힙

풀이

import sys
import heapq
input = sys.stdin.readline

n = int(input())
li = []

for _ in range(n):
    # 최대 힙 구현을 위해 -1을 곱해 역순으로 li에 들어가게 함
    x = int(input()) * -1
    if x == 0:
        if len(li) == 0:
            print(0)
        else:
            # 출력시 다시 -1을 곱해 원래 숫자로 돌려줌
            print(-heapq.heappop(li))
    else:
        heapq.heappush(li, x)