- Published on
heapq — Heap queue algorithm
- Authors
- Name
- Gene Zhang
Heaps are arrays for which a[k] <= a[2*k+1]
and a[k] <= a[2*k+2]
for all k, counting elements from 0.
The interesting property of a heap is that a[0]
is always its smallest element.
Python heapq library
Source code: Lib/heapq.py
This module provides an implementation of the heap queue algorithm, also known as the priority queue algorithm.
import heapq
li = [5, 7, 9, 1, 3]
heapq.heapify(li) # [1, 3, 9, 7, 5]
print(li)
# pushes 4 into li
heapq.heappush(li, 4) # [1, 3, 4, 7, 5, 9]
print(li)
# pop the first (smallest) element
print(heapq.heappop(li)) # 1
print(li) # [3, 5, 4, 7, 9]
# push then pop
print(heapq.heappushpop(li, 2)) # 2
print(li) # [3, 5, 4, 7, 9]
# pop then push
print(heapq.heapreplace(li, 2)) # 3
print(li) # [2, 5, 4, 7, 9]
# find the largest and smallest elements
# nlargest(k, iterable, key = fun)
# nsmallest(k, iterable, key = fun)
# Equivalent to: sorted(iterable, key=key, reverse=True)[:n].
print(heapq.nlargest(3, li)) # [9, 7, 5]
# Equivalent to: sorted(iterable, key=key)[:n].
print(heapq.nsmallest(2, li)) # [2, 4]
# A heapsort can be implemented by:
# 1. pushing all values onto a heap
# 2. popping off the smallest values one at a time
def heapsort(iterable):
h = []
for value in iterable:
heapq.heappush(h, value)
return [heapq.heappop(h) for i in range(len(h))]
print(heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])) # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
Use case: find the kth largest element in an array.
Method: build a min heap that consists of the first k largest numbers. Then the first one would be the k-th largest.
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
heap = nums[:k]
heapq.heapify(heap)
for num in nums[k:]:
# A larger element that should be included in the top k largest elements
if num > heap[0]:
heapq.heapreplace(heap, num)
return heap[0]