파이썬 heapq 모듈 사용법과 응용

✏️ heapq

최소 힙(min heap) 자료구조를 제공한다. 원소는 항상 정렬되어 추가되고 삭제되며 가장 작은 값이 언제나 인덱스 0에 위치하게 된다. 

import heapq

 

✏️ 최소힙의 생성

heapq 모듈을 이용함으로써 일반 리스트를 마치 최소 힙처럼 다룰수 있게 된다. 그냥 빈 리스트를 생성하고 heapq 모듈의 함수를 호출할 때마다 리스트를 인자로 넘겨서 사용하면된다. 즉 heapq모듈을 통해서 원소를 추가하고 삭제하면 그 리스트는 최소힙이 된다.

heap=[]

 

✏️ 최소힙 원소추가

heappush(), 첫번째 인자는 원소를 추가할 대상 리스트, 두번째 인자는 추가할 원소

시간복잡도: O(logN)

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

#output
#[1, 3, 7, 4]

 

✏️ 최소힙 원소삭제

heappop(), 대상리스트를 넘기면 가장 작은 원소를 삭제하고 값을 리턴한다.

시간복잡도: O(logN)

print(heapq.heappop(heap))
print(heap)

print(heapq.heappop(heap))
print(heap)

# output
# 1
# [3, 4, 7]
# 3
# [4, 7]

 

✏️ 기존리스트 힙으로 변환

heapift()

시간복잡도: O(N)

heap = [4, 1, 7, 3, 8, 5]
heapq.heapify(heap)
print(heap)

# output
# [1, 3, 5, 4, 8, 7]

 

✏️ 최대 힙

heapq 모듈은 최소 힙으로 동작한다고 했다. 최대 힙으로 활용하려면 어떻게 해야할까?

IDEA : y= -x 변환을 하면 최솟값 정렬이 최댓값 정렬로 바뀐다.

힙에 원소를 추가할 때 (-i, i)의 튜플형태로 넣어주면 튜플의 첫 번째 원소를 기준으로 힙을 구성하게 된다.

 

import heapq

nums = [4, 1, 7, 3, 8, 5]
heap = []

for num in nums:
  heapq.heappush(heap, (-num, num))

while heap:
  print(heapq.heappop(heap)[1])

만약 print(heap)을 해준다면 아래와 같이 출력될 것이다.

[(-8, 8), (-7, 7), (-5, 5), (-1, 1), (-3, 3), (-4, 4)]

실제값은 튜플의 두 번째 자리에 저장되어 있으니까 [1]인덱싱을 통해서 접근해주면 된다.

 

 

✏️ 백준 11279번, 최대 힙

https://www.acmicpc.net/problem/11279

11279번: 최대 힙

import heapq
from sys import stdin

N=int(stdin.readline())
heap=[]
for _ in range(N):
  x=int(stdin.readline())
  if x==0:
    if heap:
      print(heapq.heappop(heap)[1])
    else:
      print(0)
  else:
    heapq.heappush(heap, (-x, x))

* 입력시 stdin 안쓰고 input()으로 받으면 시간초과난다.