Heaps & Priority Queues

A Heap is a complete binary tree where the parent is always smaller (min-heap) or larger (max-heap) than its children. Python's heapq module is a min-heap by default.

Key operations:
• heappush(heap, val) → O(log n)
• heappop(heap)       → O(log n), returns minimum
• heap[0]             → O(1), peek minimum without removing
• heapify(list)       → O(n), convert list to heap in-place

Max-heap trick in Python: negate values. Push -val, pop returns -min = max.

When to use a heap:
• "Top K" problems → maintain heap of size K
• "K closest" points → max-heap of size K
• Merge K sorted lists → min-heap of (val, list_idx)
• Running median → two heaps (max-heap left half, min-heap right half)
• Task scheduling by priority
• Dijkstra's shortest path algorithm

1import heapq
2from collections import defaultdict
3
4# ── Basic heap operations ─────────────────────────────────────────
5heap = []
6heapq.heappush(heap, 5)
7heapq.heappush(heap, 2)
8heapq.heappush(heap, 8)
9heapq.heappush(heap, 1)
10print(heapq.heappop(heap))   # 1 (minimum)
11print(heap[0])               # 2 (peek min without removing)
12
13# Max-heap: negate values
14max_heap = []
15for val in [3, 1, 4, 1, 5, 9]:
16    heapq.heappush(max_heap, -val)
17print(-heapq.heappop(max_heap))  # 9 (maximum)
18
19# ── Top K Frequent Elements ───────────────────────────────────────
20def top_k_frequent(nums: list[int], k: int) -> list[int]:
21    count = defaultdict(int)
22    for n in nums: count[n] += 1
23    # min-heap of size k: (freq, num)
24    heap = []
25    for num, freq in count.items():
26        heapq.heappush(heap, (freq, num))
27        if len(heap) > k:
28            heapq.heappop(heap)    # removes least frequent
29    return [num for freq, num in heap]
30
31# ── K Closest Points to Origin ────────────────────────────────────
32def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
33    max_heap = []
34    for x, y in points:
35        dist = -(x*x + y*y)          # negate for max-heap
36        heapq.heappush(max_heap, (dist, x, y))
37        if len(max_heap) > k:
38            heapq.heappop(max_heap)  # removes farthest
39    return [[x, y] for _, x, y in max_heap]
40
41# ── Merge K Sorted Lists ──────────────────────────────────────────
42def merge_k_sorted(lists: list[list[int]]) -> list[int]:
43    heap   = []
44    result = []
45    for i, lst in enumerate(lists):
46        if lst:
47            heapq.heappush(heap, (lst[0], i, 0))   # (val, list_idx, elem_idx)
48    while heap:
49        val, i, j = heapq.heappop(heap)
50        result.append(val)
51        if j + 1 < len(lists[i]):
52            heapq.heappush(heap, (lists[i][j+1], i, j+1))
53    return result
54
55# ── Running Median (two heaps) ────────────────────────────────────
56class MedianFinder:
57    def __init__(self):
58        self.small = []   # max-heap (negated) for lower half
59        self.large = []   # min-heap for upper half
60
61    def add_num(self, num: int):
62        heapq.heappush(self.small, -num)
63        heapq.heappush(self.large, -heapq.heappop(self.small))
64        if len(self.large) > len(self.small):
65            heapq.heappush(self.small, -heapq.heappop(self.large))
66
67    def find_median(self) -> float:
68        if len(self.small) > len(self.large): return -self.small[0]
69        return (-self.small[0] + self.large[0]) / 2.0
70
71print(top_k_frequent([1,1,1,2,2,3], 2))             # [1, 2]
72print(k_closest([[1,3],[-2,2],[5,8],[0,1]], 2))      # [[-2,2],[0,1]]
73print(merge_k_sorted([[1,4,7],[2,5,8],[3,6,9]]))     # [1,2,3,4,5,6,7,8,9]
74mf = MedianFinder()
75for n in [1,2,3]: mf.add_num(n)
76print(mf.find_median())                              # 2.0