Published on

heapq — Heap queue algorithm

Authors
  • avatar
    Name
    Gene Zhang
    Twitter

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]