[python] heapq

파이썬의 heapq 모듈은 우선순위 큐 알고리즘의 구현을 제공한다.

따라서 heapq를 사용하여 최소힙과 최대힙을 구현할 수 있다.

 

heappush(heap, item)

힙의 조건을 유지하며 item을 heap으로 push한다.

heappop(heap)

힙의 조건을 유지하며 heap에서 가장 작은 item을 pop하고 반환한다.

heap이 비어있다면 IndexError가 발생한다.

 

 

최소 힙

heapq는 기본적으로 최소 힙을 구현한다.

백준에 최소 힙을 구현하는 문제가 있어 적용하여 보았다.

1927번: 최소 힙 (acmicpc.net)

1927번: 최소 힙

import heapq
import sys

input = sys.stdin.readline

lst =[]

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

 

 

최대 힙

heapq는 기본적으로 최소 힙을 구현하므로 최대 힙을 구하기 위해서는 약간의 트릭을 사용하면 된다.

마찬가지로 백준에 최대 힙을 구현하는 문제가 있어 적용하여 보았다.

11279번: 최대 힙 (acmicpc.net)

11279번: 최대 힙

 

import heapq
import sys
input = sys.stdin.readline

lst=[]
for _ in range(int(input())):
    x = int(input())
    
    if x != 0:
        heapq.heappush(lst, -x)
    else:
        if len(lst) == 0:
            print(0)
        else:
            y = heapq.heappop(lst)
            print(-y)

 

저장하는 값과 pop한 값에 -를 붙여줌으로써 최대 힙을 구현한다.