서론

파이썬 우선순위큐 라이브러리중 heapq 를 사용하며 알게 된 점을 정리하고자 한다.

heapq 라이브러리

heapq 는 우선순위 큐 알고리즘이라고도 하는 힙 큐 알고리즘의 구현을 제공한다.

힙은 모든 부모 노드가 자식보다 작거나 같은 값을 갖는 이진 트리이다.
힙의 흥미로운 특성은 가장 작은 요소가 항상 루트인 heap[0] 이라는 것이다.

heap 을 만들려면 [] 로 초기화된 리스트를 사용하거나 함수 heapify() 를 통해 값이 들어 있는 리스트를 힙으로 변환할 수 있다.

heapq 메서드

heapq.heappush(heap, item)

힙 불변성을 유지하면서 item 값을 heap 으로 푸시한다.

que = []
heapq.heappush(que, 3)

que = []
heapq.heappush(que, (0, 2))

heapq.heappop(heap)

힙 불변성을 유지하면서 heap 에서 가장 작은 항목을 팝하고 반환한다. 힙이 비어 있으면 IndexError 가 발생한다. 팝하지 않고 가장 작은 항목에 액세스하려면 heap[0] 을 사용하자.

que = []
heapq.heappush(que, 3)
heapq.heappush(que, 1)
heapq.heappush(que, 2)

print(heapq.heappop(que)) # 1
print(heapq.heappop(que)) # 2
print(heapq.heappop(que)) # 3

heapq.heapify(x)

위에서 힙에 값을 push 할때 값 하나하나 넣지 않고 기존 배열을 힙으로 변환시켜줄 수 있다.

que = []
heapq.heappush(que, 3)
heapq.heappush(que, 1)
heapq.heappush(que, 2)

# 위와 동일
nums = [3,1,2]
heapq.heapify(nums)

heapq.nlargest(n, iterable, key=None)

iterable 에 의해 정의된 데이터 집합에서 n 개의 가장 큰 요소로 구성된 리스트를 반환한다. key 가 제공되면 iteralbe 의 각 요소에서 비교 키를 추출하는 데 사용되는 단일 인자 함수를 지정한다.

sorted(iterable, key=key, reverse=True)[:n] 과 동일하다.

nums = [3,1,2]
heapq.heapify(nums)

heapq.nlargest(2, nums) # [3, 2]

leetcode 215. Kth Largest Element in an Array 문제

https://leetcode.com/problems/kth-largest-element-in-an-array/description/

풀이코드

import heapq
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        nums = [-1 * x for x in nums]
        heapq.heapify(nums)
        for _ in range(k - 1):
            heapq.heappop(nums)
        return -heapq.heappop(nums)

공간복잡도를 낮추기 위해 따로 배열을 할당하지 않고 부호만 바꾸고 heappop() 을 사용해서 큰 순서대로 값을 뽑아줬다.