Heap / Priority Queue

Use when: K largest/smallest, streaming median, greedy problems needing the next best element.

Complexity: O(log N) push/pop, O(N) heapify

Min Heap (default in Python)

import heapq
 
heap = []
heapq.heappush(heap, val)
smallest = heapq.heappop(heap)
peek = heap[0]
 
# Heapify in O(N)
heapq.heapify(nums)

Max Heap (negate values)

import heapq
 
heap = []
heapq.heappush(heap, -val)
largest = -heapq.heappop(heap)

K Largest Elements

import heapq
 
def k_largest(nums, k):
    return heapq.nlargest(k, nums)
 
# Or with min heap of size k:
def k_largest_heap(nums, k):
    heap = nums[:k]
    heapq.heapify(heap)
    for n in nums[k:]:
        if n > heap[0]:
            heapq.heapreplace(heap, n)
    return heap

K Smallest Elements

import heapq
 
def k_smallest(nums, k):
    return heapq.nsmallest(k, nums)

Kth Largest

import heapq
 
def find_kth_largest(nums, k):
    heap = nums[:k]
    heapq.heapify(heap)
    for n in nums[k:]:
        if n > heap[0]:
            heapq.heapreplace(heap, n)
    return heap[0]

Heap with Custom Key (tuples)

# Push tuples - compared lexicographically
heapq.heappush(heap, (priority, item))
priority, item = heapq.heappop(heap)

Signals in a Problem

  • “K largest/smallest/closest”
  • “Median of data stream”
  • Greedy: always pick the best available option
  • Merge K sorted lists