8. Heaps

What Interviewers Test Here

• heapq for streaming top-K queries
• Two-heap pattern for median maintenance
• When to prefer a heap over sorting

import heapq

# --- K largest elements --- O(n log k)
def k_largest(nums: list[int], k: int) -> list[int]:
    return heapq.nlargest(k, nums)

# --- K closest to origin ---
import math
def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
    return heapq.nsmallest(k, points, key=lambda p: p[0]**2 + p[1]**2)

# --- Merge k sorted lists ---
def merge_k_sorted(lists) -> list[int]:
    heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
    result = []
    while heap:
        val, i, j = heapq.heappop(heap)
        result.append(val)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j+1], i, j+1))
    return result

---

Conceptual Q&A

Q1: "Find the kth largest element in an unsorted array."

Walkthrough:

"Several approaches:

1. Sort — $O(n \log n)$, then return nums[-k]. Simple but slower.
2. Min-heap of size k — $O(n \log k)$. Maintain a heap of the k largest seen so far. Heap root is always the kth largest. When a new element is larger than the root, pop and push.
3. Quickselect — $O(n)$ average. Like quicksort, but only recurse into the side containing the kth element.

In an interview: state all three, code #2 (most reliable). Mention #3 for the follow-up 'can you do better?'"

def kth_largest(nums: list[int], k: int) -> int:
    # Min-heap of size k: root = kth largest
    heap = nums[:k]
    heapq.heapify(heap)                  # O(k)
    for x in nums[k:]:                   # O((n-k) log k)
        if x > heap[0]:
            heapq.heapreplace(heap, x)   # pop min, push x — O(log k)
    return heap[0]

# One-liner alternative (less efficient for large n):
# return heapq.nlargest(k, nums)[-1]

---

Q2: "How would you find the running median of a stream of numbers?"

Walkthrough:

"The median depends on the sorted order, which changes as numbers arrive. Sorting after each addition is $O(n \log n)$ per query.

Better: maintain two heaps that split the data at the median:

• Lower half (max-heap, negated): stores all numbers ≤ median
• Upper half (min-heap): stores all numbers ≥ median

Keep them balanced: lower can have at most 1 more element than upper. Median is: if unequal sizes → top of larger heap; if equal → average of both tops.

Add in $O(\log n)$, query in $O(1)$."

class MedianFinder:
    def __init__(self):
        self.lower: list[int] = []   # max-heap (negate)
        self.upper: list[int] = []   # min-heap

    def add_num(self, num: int) -> None:
        heapq.heappush(self.lower, -num)
        # Always push max of lower to upper first (maintain ordering)
        heapq.heappush(self.upper, -heapq.heappop(self.lower))
        # Rebalance: lower can be equal size or one larger
        if len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def find_median(self) -> float:
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])
        return (-self.lower[0] + self.upper[0]) / 2.0

---