Vasu Patel logoVasu Patel
← Python DSA — Course
All chapters in this track
  1. 011. Complexity Analysis
  2. 022. Arrays & Strings
  3. 033. Hashing
  4. 044. Stacks & Queues
  5. 055. Linked Lists
  6. 066. Binary Search
  7. 077. Sorting
  8. 088. Recursion & Backtracking
  9. 099. Trees
  10. 1010. Heaps
  11. 1111. Graphs
  12. 1212. Dynamic Programming
  13. 1313. Tries
  14. 1414. Advanced Structures
  15. 15Appendix: CP Quick Reference

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

OperationTime
heappushO(log⁡n)O(\log n)
heappopO(log⁡n)O(\log n)
heapifyO(n)O(n)
nlargest(k, it)O(n+klog⁡n)O(n + k \log n)
nsmallest(k, it)O(n+klog⁡n)O(n + k \log n)

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]