S029 · Heaps & Priority Queues

"""heaps.py — the interview canon for heaps + a mini Dijkstra.
 
Run:  python heaps.py
"""
from __future__ import annotations
import heapq
from collections import Counter
from typing import Any
 
 
# ------------------------------------------------------------------
# 1) Top-K frequent — min-heap of size K
# ------------------------------------------------------------------
def top_k_frequent(nums: list[int], k: int) -> list[int]:
    """Return the K most frequent values. O(n log k) time, O(k) heap memory."""
    counts = Counter(nums)
    heap: list[tuple[int, int]] = []            # (freq, value); freq is the sort key
    for val, freq in counts.items():
        heapq.heappush(heap, (freq, val))
        if len(heap) > k:
            heapq.heappop(heap)                 # drop the least-frequent
    return [v for _, v in heap]
 
 
# ------------------------------------------------------------------
# 2) K-th largest — same pattern, single number out
# ------------------------------------------------------------------
def kth_largest(nums: list[int], k: int) -> int:
    heap: list[int] = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap[0]                              # k-th largest = smallest in the heap of top K
 
 
# ------------------------------------------------------------------
# 3) Streaming median — two heaps trick
# ------------------------------------------------------------------
class MedianFinder:
    """Add numbers one at a time, get median in O(log n).
 
    Trick: two heaps.
      - lower: max-heap of the smaller half (stored as negated values in Python's min-heap)
      - upper: min-heap of the larger half
    Invariant: len(lower) == len(upper) or len(lower) == len(upper) + 1.
    """
    def __init__(self) -> None:
        self.lower: list[int] = []              # max-heap (negated)
        self.upper: list[int] = []              # min-heap
 
    def add(self, num: int) -> None:
        # Push into the correct half
        if not self.lower or num <= -self.lower[0]:
            heapq.heappush(self.lower, -num)
        else:
            heapq.heappush(self.upper, num)
        # Rebalance so sizes differ by at most 1, and lower may be exactly one larger
        if len(self.lower) > len(self.upper) + 1:
            heapq.heappush(self.upper, -heapq.heappop(self.lower))
        elif len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))
 
    def median(self) -> float:
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])
        return (-self.lower[0] + self.upper[0]) / 2
 
 
# ------------------------------------------------------------------
# 4) Merge K sorted lists — heap of head pointers
# ------------------------------------------------------------------
def merge_k_sorted(lists: list[list[int]]) -> list[int]:
    """O(n log k) where n = total elements, k = number of lists."""
    heap: list[tuple[int, int, int]] = []       # (value, list_idx, item_idx)
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
 
    out: list[int] = []
    while heap:
        val, i, j = heapq.heappop(heap)
        out.append(val)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
    return out
 
 
# ------------------------------------------------------------------
# 5) Dijkstra's shortest path — heap of (dist, node) with lazy deletion
# ------------------------------------------------------------------
def dijkstra(graph: dict[int, list[tuple[int, int]]], start: int) -> dict[int, float]:
    """graph: node -> list of (neighbor, edge_weight). Returns dist[node] for all reachable."""
    dist: dict[int, float] = {start: 0}
    heap: list[tuple[float, int]] = [(0.0, start)]
    while heap:
        d, u = heapq.heappop(heap)
        if d > dist.get(u, float("inf")):
            continue                            # stale entry — skip
        for v, w in graph.get(u, []):
            nd = d + w
            if nd < dist.get(v, float("inf")):
                dist[v] = nd
                heapq.heappush(heap, (nd, v))
    return dist
 
 
# ------------------------------------------------------------------
# 6) Event scheduler — priority queue by run-at timestamp
# ------------------------------------------------------------------
class Scheduler:
    """Minimal cron-like scheduler. Heap of (run_at, id, callback)."""
    def __init__(self) -> None:
        self._heap: list[tuple[float, int, Any]] = []
        self._next_id = 0
 
    def schedule(self, run_at: float, callback) -> None:
        heapq.heappush(self._heap, (run_at, self._next_id, callback))
        self._next_id += 1                      # tie-breaker to keep tuples comparable
 
    def next_due(self, now: float):
        """Return the next callback due by ‘now’, else None."""
        if self._heap and self._heap[0][0] <= now:
            return heapq.heappop(self._heap)[2]
        return None
 
 
if __name__ == "__main__":
    print("top_k_frequent [1,1,1,2,2,3] k=2 :", top_k_frequent([1, 1, 1, 2, 2, 3], 2))
    print("kth_largest [3,2,1,5,6,4] k=2 :", kth_largest([3, 2, 1, 5, 6, 4], 2))
 
    mf = MedianFinder()
    for x in [1, 5, 2, 10, 7]:
        mf.add(x)
    print("median after [1,5,2,10,7]:", mf.median())     # 5
 
    lists = [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
    print("merge_k_sorted :", merge_k_sorted(lists))
 
    graph = {
        1: [(2, 7), (3, 9), (6, 14)],
        2: [(3, 10), (4, 15)],
        3: [(4, 11), (6, 2)],
        4: [(5, 6)],
        5: [],
        6: [(5, 9)],
    }
    print("dijkstra from 1 :", dijkstra(graph, 1))       # {1:0, 2:7, 3:9, 4:20, 6:11, 5:20}
 
    sched = Scheduler()
    sched.schedule(100, lambda: "task A")
    sched.schedule(50, lambda: "task B")
    print("next_due at t=60 :", sched.next_due(60)())     # task B (earlier)