Back to Notes

Heap & Priority Queue

🧠 Spot It

  • "Top K", "Kth largest/smallest", "K most frequent"
  • Need to repeatedly get the min or max efficiently
  • Merging multiple sorted streams

Python's heapq is a min-heap by default. For max-heap: push -value.


Python heapq Basics

import heapq

# Min-heap
heap = []
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 2)
heapq.heappop(heap)   # returns 1 (smallest)
heap[0]               # peek without removing

# Max-heap trick: negate values
heapq.heappush(heap, -val)
max_val = -heapq.heappop(heap)

# Heapify existing list in O(N)
heapq.heapify(arr)

Pattern 1 — Top K Elements

def top_k_frequent(nums, k):
    count = {}
    for n in nums:
        count[n] = count.get(n, 0) + 1
    # Min-heap of size K — keep K largest
    heap = []
    for num, freq in count.items():
        heapq.heappush(heap, (freq, num))
        if len(heap) > k:
            heapq.heappop(heap)  # remove smallest freq
    return [item[1] for item in heap]

Pattern 2 — Two Heaps (Median)

# max_heap (left half) + min_heap (right half)
# median = top of max_heap or avg of both tops
import heapq

class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (negated)
        self.large = []  # min-heap

    def addNum(self, num):
        heapq.heappush(self.small, -num)
        # Balance: ensure small's max <= large's min
        if self.small and self.large and -self.small[0] > self.large[0]:
            heapq.heappush(self.large, -heapq.heappop(self.small))
        # Balance sizes
        if len(self.small) > len(self.large) + 1:
            heapq.heappush(self.large, -heapq.heappop(self.small))
        if len(self.large) > len(self.small):
            heapq.heappush(self.small, -heapq.heappop(self.large))

    def findMedian(self):
        if len(self.small) > len(self.large):
            return -self.small[0]
        return (-self.small[0] + self.large[0]) / 2

Problems

#ProblemPatternLinkStatus
1Kth Largest Element in ArrayTop KLC 215
2Top K Frequent ElementsTop K + freq mapLC 347
3Find Median from Data StreamTwo HeapsLC 295
4K Closest Points to OriginTop KLC 973
5Task SchedulerHeap + greedyLC 621
6Design TwitterHeap + K-way mergeLC 355

Notes

<!-- Add your own observations as you solve problems -->