👀 heapq와 우선순위 큐
파이썬의 heapq는 최소 힙(min heap)을 구현할 때 사용할 수 있는 모듈로, 우선순위 큐(priority queue)를 구현할 때 많이 사용된다.
아래 그림은 최대 힙의 예시로 루트 노드(배열에서는 0번째 요소)에 최대값이 위치한다. 최소 힙은 반대로 루트노드에 최소값이 위치한다.
💻 heapq 주요 함수
📌 heapq.heapify(x)
주어진 리스트 x를 힙으로 변환. 시간복잡도는 O(n).
📌 heapq.heappush(heap, item)
주어진 힙(heap)에 원소(item)를 삽입. 시간복잡도는 O(log n).
최대 힙을 구현하고자 할 때에는 item에 음의 부호를 붙여서 넣어주면 된다.
📌 heapq.heappop(heap)
힙에서 최소 원소를 반환하면서 삭제. 시간복잡도는 O(log n).
📝 heapq를 사용하여 정렬하기
리스트를 heapify를 사용하여 힙으로 변환하고, heappop를 사용하여 차례대로 꺼내 다시 리스트에 담으면 오름차순으로 정렬된 리스트를 얻을 수 있다.
import heapq
def heapsort(iterable):
heap = list(iterable)
heapq.heapify(heap) # 입력 리스트를 힙으로 변환
result = []
while heap:
result.append(heapq.heappop(heap))
return result📚 heapq를 사용해서 풀 수 있는 문제 목록
- 백준 13334번 철로
- 백준 2750번 수 정렬하기
- 백준 11279번 최대 힙
- 백준 1432번 그래프 수정
- 백준 1781번 컵라면
- 백준 1655번 가운데를 말해요
- 백준 3055번 탈출
- 백준 2665번 미로
- 백준 1715번 카드 정렬하기
- LeetCode 215. Kth Largest Element in an Array
💚 예시: 백준 1655번 가운데를 말해요
이 문제는 주어진 수들을 정렬했을 때 가장 가운데에 있는 수를 찾는 문제이다.
주어진 수가 짝수 개라면 가운데에 있는 수 중 작은 수를 반환해야 한다.
이 문제는 두 개의 힙을 활용해서 풀 수 있다.
첫 번째 힙 low는 전체 값들 중 작은 값 절반을 가지도록 하고, max heap으로 구현해서 최상단에 위치한 값은 작은 값들 중 가장 큰 값이 되도록 한다.
그리고 두 번째 힙 high는 전체 값들 중 큰 값 절반을 가지도록 하고, min heap으로 구현해서 최상단에 위치한 값은 큰 값들 중 가장 작은 값이 되도록 한다.
주어진 수가 짝수 개라면 가운데에 있는 수 중 작은 수를 반환해야 하므로, low 힙의 루트에 위치한 값이 답이 된다.
구체적으로 어떻게 구현하면 되는지는 아래 코드와 주석을 통해 살펴보자.
import sys
import heapq
n = int(sys.stdin.readline())
low = []
high = []
for i in range(n):
value = int(sys.stdin.readline())
# 새로 추가된 값은 low 힙에 추가. low는 최대 힙이므로 -를 붙임
heapq.heappush(low, -value)
# low 힙에서 가장 큰 값을 high 힙으로 보내기
heapq.heappush(high, -heapq.heappop(low))
# 만약 전체 요소의 개수가 홀수라면 low 힙이 요소를 하나 더 가지도록 함
if len(high) > len(low):
heapq.heappush(low, -heapq.heappop(high))
# 가운데 값 출력
print(-low[0])