All chapters in this track
Chapter 10 of 15
10. Heaps
Key Insight
Python's heapq is a min-heap. For max-heap, negate values. Master: streaming k-th largest/smallest, merge k sorted lists, and lazy deletion (mark-invalid items instead of removing them).
Complexity
| Operation | Time |
|---|---|
heappush | |
heappop | |
heapify | |
nlargest(k, it) | |
nsmallest(k, it) |
10.1 Core Templates
import heapq
# --- Min-heap basics ---
h = []
heapq.heappush(h, 3)
heapq.heappush(h, 1)
heapq.heappush(h, 2)
print(heapq.heappop(h)) # 1 — always smallest
print(h[0]) # peek without removing
# --- Max-heap (negate values) ---
max_h = []
for x in [3, 1, 4, 1, 5]:
heapq.heappush(max_h, -x)
print(-heapq.heappop(max_h)) # 5
# --- Heapify in O(n) ---
arr = [3, 1, 4, 1, 5, 9]
heapq.heapify(arr) # in-place
# --- Heap of tuples: (priority, value) ---
tasks = []
heapq.heappush(tasks, (1, "low priority"))
heapq.heappush(tasks, (0, "high priority"))
pri, task = heapq.heappop(tasks) # (0, "high priority")
10.2 K-th Largest / Smallest
# --- Kth largest in a stream — O(n log k) ---
class KthLargest:
def __init__(self, k: int, nums: list[int]):
self.k = k
self.heap: list[int] = []
for x in nums:
self.add(x)
def add(self, val: int) -> int:
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]
# --- Kth largest in array — O(n + k log n) ---
def kth_largest(nums: list[int], k: int) -> int:
heapq.heapify(nums) # min-heap
while len(nums) > k:
heapq.heappop(nums)
return nums[0]
# Alternative one-liner: return heapq.nlargest(k, nums)[-1]
# --- Median from data stream ---
class MedianFinder:
def __init__(self):
self.lo: list[int] = [] # max-heap (negated) — lower half
self.hi: list[int] = [] # min-heap — upper half
def add_num(self, num: int) -> None:
heapq.heappush(self.lo, -num)
heapq.heappush(self.hi, -heapq.heappop(self.lo))
if len(self.hi) > len(self.lo):
heapq.heappush(self.lo, -heapq.heappop(self.hi))
def find_median(self) -> float:
if len(self.lo) > len(self.hi):
return float(-self.lo[0])
return (-self.lo[0] + self.hi[0]) / 2
10.3 Merge K Sorted Lists / Arrays
# --- Merge k sorted lists --- O(n log k)
def merge_k_lists(lists: list[ListNode | None]) -> ListNode | None:
heap = []
for i, node in enumerate(lists):
if node:
heapq.heappush(heap, (node.val, i, node))
dummy = ListNode()
cur = dummy
while heap:
val, i, node = heapq.heappop(heap)
cur.next = node; cur = cur.next
if node.next:
heapq.heappush(heap, (node.next.val, i, node.next))
return dummy.next
# --- Smallest range covering at least one from each of k lists --- O(n log k)
def smallest_range(nums: list[list[int]]) -> list[int]:
heap = [(row[0], i, 0) for i, row in enumerate(nums)]
heapq.heapify(heap)
cur_max = max(row[0] for row in nums)
best = [heap[0][0], cur_max]
while heap:
cur_min, i, j = heapq.heappop(heap)
if cur_max - cur_min < best[1] - best[0]:
best = [cur_min, cur_max]
if j + 1 == len(nums[i]):
break
next_val = nums[i][j+1]
heapq.heappush(heap, (next_val, i, j+1))
cur_max = max(cur_max, next_val)
return best
10.4 Lazy Deletion
# When you can't efficiently remove arbitrary elements from a heap,
# push a "tombstone" and skip it on pop.
import heapq
from collections import defaultdict
class LazyHeap:
def __init__(self):
self._heap: list = []
self._removed: dict = defaultdict(int)
def push(self, val):
heapq.heappush(self._heap, val)
def remove(self, val):
self._removed[val] += 1 # mark for lazy removal
def pop(self):
while self._heap and self._removed[self._heap[0]] > 0:
val = heapq.heappop(self._heap)
self._removed[val] -= 1
return heapq.heappop(self._heap)
def peek(self):
while self._heap and self._removed[self._heap[0]] > 0:
val = heapq.heappop(self._heap)
self._removed[val] -= 1
return self._heap[0]