Heaps & Priority Queues, Part 1: The Structure Behind Top-K

# Heap stored as an array:
#
#           1           index 0
#          / \
#         3   5         index 1, 2
#        / \ / \
#       7  8 6  9       index 3, 4, 5, 6
#
# Array: [1, 3, 5, 7, 8, 6, 9]
#
# Parent of index 3 (value 7)? (3 - 1) // 2 = 1  (value 3) ✓
# Parent of index 5 (value 6)? (5 - 1) // 2 = 2  (value 5) ✓
# Left child of index 1 (value 3)? 2*1+1 = 3     (value 7) ✓
# Right child of index 1 (value 3)? 2*1+2 = 4    (value 8) ✓

heap = [1, 3, 5, 7, 8, 6, 9]

def parent(i: int) -> int:
    return (i - 1) // 2

def left_child(i: int) -> int:
    return 2 * i + 1

def right_child(i: int) -> int:
    return 2 * i + 2

# Verify heap property: every parent <= its children
def is_min_heap(arr: list) -> bool:
    """Check if array satisfies the min-heap property."""
    n = len(arr)
    for i in range(n):
        l = left_child(i)
        r = right_child(i)
        if l < n and arr[i] > arr[l]:
            return False
        if r < n and arr[i] > arr[r]:
            return False
    return True

print(is_min_heap([1, 3, 5, 7, 8, 6, 9]))   # True
print(is_min_heap([1, 3, 5, 2, 8, 6, 9]))   # False (2 < parent 3, but at index 3, parent is index 1 = 3, 2 < 3)

import heapq

# --- heapq.heappush(h, item) --- O(log n)
# Insert an item while maintaining the heap property
h = []
heapq.heappush(h, 5)
heapq.heappush(h, 1)
heapq.heappush(h, 8)
heapq.heappush(h, 3)
print(h)  # [1, 3, 8, 5]  -- internal array (NOT sorted, just heap-valid)

# --- heapq.heappop(h) --- O(log n)
# Remove and return the minimum element
print(heapq.heappop(h))  # 1  (the minimum)
print(heapq.heappop(h))  # 3
print(h)                 # [5, 8]

# --- heapq.heapify(list) --- O(n)  in-place conversion (Floyd's algorithm)
data = [10, 4, 7, 1, 9, 2, 6]
heapq.heapify(data)
print(data)  # [1, 4, 2, 10, 9, 7, 6]  -- heap-valid min-heap

# --- heapq.heappushpop(h, item) --- O(log n)
# Push item, then pop and return the minimum. More efficient than push+pop.
h = [1, 3, 5]
heapq.heapify(h)
result = heapq.heappushpop(h, 2)
print(result)  # 1  (pushed 2, then popped the new min which is 1)
print(h)       # [2, 3, 5]

# --- heapq.heapreplace(h, item) --- O(log n)
# Pop and return the minimum, then push item. Raises IndexError if empty.
# Faster than heappop + heappush because it avoids one sift-up.
h = [1, 3, 5]
heapq.heapify(h)
result = heapq.heapreplace(h, 4)
print(result)  # 1  (popped min, then pushed 4)
print(h)       # [3, 4, 5]

# --- heapq.nlargest(n, iterable, key=None) --- O(n log k)
# Return the k largest items (more efficient than sorted()[-k:] for small k)
scores = [88, 95, 72, 100, 63, 91, 77, 85]
print(heapq.nlargest(3, scores))         # [100, 95, 91]

students = [("Alice", 88), ("Bob", 95), ("Carol", 72), ("Dana", 100)]
print(heapq.nlargest(2, students, key=lambda s: s[1]))
# [('Dana', 100), ('Bob', 95)]

# --- heapq.nsmallest(n, iterable, key=None) --- O(n log k)
print(heapq.nsmallest(3, scores))        # [63, 72, 77]

# When to prefer nlargest/nsmallest over sorted():
# - If k << n: nlargest is O(n log k) vs sorted's O(n log n)
# - If k is close to n: just use sorted() -- it is faster overall

import heapq

# Max-heap via negation
max_heap = []
for value in [5, 1, 8, 3, 9, 2]:
    heapq.heappush(max_heap, -value)  # store negated

print(max_heap)  # [-9, -8, -5, -1, -3, -2] -- internally a min-heap of negatives

# Extract maximum: pop and negate
while max_heap:
    print(-heapq.heappop(max_heap), end=" ")
# 9 8 5 3 2 1  -- extracted in descending order

# Priority queue pattern: (priority, item) tuples
# Lower number = higher priority (min-heap pops smallest first)
task_queue = []
heapq.heappush(task_queue, (3, "low-priority task"))
heapq.heappush(task_queue, (1, "urgent task"))
heapq.heappush(task_queue, (2, "medium-priority task"))

while task_queue:
    priority, task = heapq.heappop(task_queue)
    print(f"Processing (priority {priority}): {task}")
# Processing (priority 1): urgent task
# Processing (priority 2): medium-priority task
# Processing (priority 3): low-priority task

# Tie-breaking with a counter for stable ordering
import itertools
counter = itertools.count()  # monotonically increasing counter

stable_queue = []
heapq.heappush(stable_queue, (2, next(counter), "first medium task"))
heapq.heappush(stable_queue, (2, next(counter), "second medium task"))
# Same priority? The counter ensures FIFO order between equal priorities.

class MinHeap:
    """A min-heap implemented over a Python list."""

    def __init__(self) -> None:
        self._data: list = []

    # --- Index helpers ---
    def _parent(self, i: int) -> int:
        return (i - 1) // 2

    def _left(self, i: int) -> int:
        return 2 * i + 1

    def _right(self, i: int) -> int:
        return 2 * i + 2

    # --- Core operations ---
    def push(self, val) -> None:
        """Insert val into the heap. O(log n)."""
        self._data.append(val)        # 1. Add at the end
        self._sift_up(len(self._data) - 1)  # 2. Restore heap property upward

    def pop(self) -> int:
        """Remove and return the minimum. O(log n)."""
        if not self._data:
            raise IndexError("Pop from empty heap")
        # Swap root with last element
        self._data[0], self._data[-1] = self._data[-1], self._data[0]
        minimum = self._data.pop()    # remove the old root (now at end)
        if self._data:
            self._sift_down(0)        # restore heap property downward
        return minimum

    def peek(self):
        """Return the minimum without removing it. O(1)."""
        if not self._data:
            raise IndexError("Peek at empty heap")
        return self._data[0]

    def __len__(self) -> int:
        return len(self._data)

    # --- Heap property restoration ---
    def _sift_up(self, i: int) -> None:
        """Bubble element at index i upward until heap property is restored.

        Visual: newly inserted element at the bottom competes with its parent.
        If it is smaller, they swap. This repeats until the element reaches
        its correct position or the root.

        Before: [1, 3, 5, 7, 8, 6, 9, 2]  (2 just inserted at index 7)
                                        ^
        Step 1: parent(7)=3 has value 7. 2 < 7 → swap
                [1, 3, 5, 2, 8, 6, 9, 7]
        Step 2: parent(3)=1 has value 3. 2 < 3 → swap
                [1, 2, 5, 3, 8, 6, 9, 7]
        Step 3: parent(1)=0 has value 1. 2 > 1 → stop
        """
        while i > 0:
            p = self._parent(i)
            if self._data[i] < self._data[p]:
                self._data[i], self._data[p] = self._data[p], self._data[i]
                i = p
            else:
                break

    def _sift_down(self, i: int) -> None:
        """Bubble element at index i downward until heap property is restored.

        Visual: after removing the root, the last element is placed at the top.
        It then competes with its smaller child, swapping downward until correct.

        Before: [9, 3, 5, 7, 8, 6]  (9 was placed at root after pop)
        Step 1: children of 0 are index 1 (value 3) and index 2 (value 5).
                Smaller child is 3. 9 > 3 → swap.
                [3, 9, 5, 7, 8, 6]
        Step 2: children of 1 are index 3 (value 7) and index 4 (value 8).
                Smaller child is 7. 9 > 7 → swap.
                [3, 7, 5, 9, 8, 6]
        Step 3: children of 3 are index 7, 8 -- out of bounds. Stop.
        """
        n = len(self._data)
        while True:
            smallest = i
            l = self._left(i)
            r = self._right(i)
            if l < n and self._data[l] < self._data[smallest]:
                smallest = l
            if r < n and self._data[r] < self._data[smallest]:
                smallest = r
            if smallest == i:
                break  # heap property satisfied
            self._data[i], self._data[smallest] = self._data[smallest], self._data[i]
            i = smallest

    def heapify(self, arr: list) -> None:
        """Build a heap from an unordered list in O(n) using Floyd's algorithm.

        Key insight: leaf nodes (indices n//2 to n-1) are already valid heaps
        of size 1. We only need to sift_down from the last internal node upward.
        This is O(n), NOT O(n log n) — most nodes are near the bottom and
        travel very short distances.
        """
        self._data = arr.copy()
        n = len(self._data)
        for i in range(n // 2 - 1, -1, -1):  # start from last internal node
            self._sift_down(i)

    def __repr__(self) -> str:
        return f"MinHeap({self._data})"

# --- Demo ---
h = MinHeap()
for v in [5, 3, 8, 1, 9, 2, 6]:
    h.push(v)

print(h)              # MinHeap([1, 3, 2, 5, 9, 8, 6])
print(h.peek())       # 1
print(h.pop())        # 1
print(h.pop())        # 2
print(h)              # MinHeap([3, 5, 6, 8, 9])

# Floyd's heapify -- O(n)
h2 = MinHeap()
h2.heapify([10, 4, 7, 1, 9, 2, 6])
print(h2)             # MinHeap([1, 4, 2, 10, 9, 7, 6])

# Extract sorted order (heap sort)
sorted_output = []
while len(h2):
    sorted_output.append(h2.pop())
print(sorted_output)  # [1, 2, 4, 6, 7, 9, 10]

import heapq
from dataclasses import dataclass, field
from typing import Any

# --- Pattern 1: Simple tuple-based priority queue ---
pq: list = []
heapq.heappush(pq, (1, "critical bug fix"))
heapq.heappush(pq, (3, "refactor module"))
heapq.heappush(pq, (2, "add new feature"))
heapq.heappush(pq, (1, "security patch"))  # same priority as critical bug

while pq:
    p, task = heapq.heappop(pq)
    print(f"[P{p}] {task}")
# [P1] critical bug fix
# [P1] security patch
# [P2] add new feature
# [P3] refactor module

# --- Pattern 2: Stable ordering with counter ---
import itertools
_counter = itertools.count()

def push_task(pq, priority, item):
    """Insert with tie-breaking counter for FIFO order among equal priorities."""
    heapq.heappush(pq, (priority, next(_counter), item))

# --- Pattern 3: queue.PriorityQueue (thread-safe wrapper over heapq) ---
from queue import PriorityQueue

thread_safe_pq = PriorityQueue()
thread_safe_pq.put((2, "task B"))
thread_safe_pq.put((1, "task A"))
thread_safe_pq.put((3, "task C"))

while not thread_safe_pq.empty():
    print(thread_safe_pq.get())
# (1, 'task A')
# (2, 'task B')
# (3, 'task C')

# --- Pattern 4: Max-priority queue (highest number = most important) ---
max_pq: list = []
for priority, task in [(3, "low"), (10, "urgent"), (7, "medium")]:
    heapq.heappush(max_pq, (-priority, task))  # negate priority

while max_pq:
    neg_p, task = heapq.heappop(max_pq)
    print(f"[P{-neg_p}] {task}")
# [P10] urgent
# [P7]  medium
# [P3]  low

import heapq

def kth_largest(nums: list[int], k: int) -> int:
    """Find the k-th largest element using a min-heap of size k.

    Strategy: maintain a min-heap of the k largest elements seen so far.
    The root of the heap is always the k-th largest.

    Why min-heap of size k? When the heap has k elements, its minimum
    (the root) is the k-th largest. Any new element larger than the root
    replaces it, maintaining the invariant.
    """
    min_heap: list[int] = []
    for num in nums:
        heapq.heappush(min_heap, num)
        if len(min_heap) > k:
            heapq.heappop(min_heap)  # discard smallest -- not in top-k
    return min_heap[0]  # root = k-th largest

print(kth_largest([3, 2, 1, 5, 6, 4], k=2))   # 5
print(kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], k=4))  # 4

import heapq

def merge_k_sorted(lists: list[list[int]]) -> list[int]:
    """Merge k sorted lists into one sorted list.

    Strategy: use a min-heap with (value, list_index, element_index).
    Always extract the global minimum efficiently.

    Time: O(n log k) where n = total elements, k = number of lists.
    """
    result: list[int] = []
    heap: list[tuple] = []

    # Initialize heap with the first element from each non-empty list
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))

    while heap:
        val, list_idx, elem_idx = heapq.heappop(heap)
        result.append(val)
        # Push the next element from the same list
        next_idx = elem_idx + 1
        if next_idx < len(lists[list_idx]):
            heapq.heappush(heap, (lists[list_idx][next_idx], list_idx, next_idx))

    return result

lists = [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
print(merge_k_sorted(lists))  # [1, 2, 3, 4, 5, 6, 7, 8, 9]

import heapq
from collections import Counter

def top_k_frequent(nums: list[int], k: int) -> list[int]:
    """Return the k most frequent elements.

    Step 1: count frequencies with Counter — O(n)
    Step 2: use nlargest to find top-k — O(n log k)
    """
    counts = Counter(nums)
    return heapq.nlargest(k, counts, key=counts.get)

print(top_k_frequent([1, 1, 1, 2, 2, 3], k=2))     # [1, 2]
print(top_k_frequent([1, 2], k=2))                  # [1, 2]

import heapq

class MedianFinder:
    """Maintain a running median using two heaps.

    Strategy: split the stream into two halves.
    - lower_max_heap: max-heap of the lower half (stored negated)
    - upper_min_heap: min-heap of the upper half

    Invariant: lower_max_heap.size == upper_min_heap.size
               OR lower_max_heap.size == upper_min_heap.size + 1

    Median: if sizes equal -> average of both tops
            if lower is bigger -> lower's top
    """

    def __init__(self) -> None:
        self._lower: list[int] = []  # max-heap (negated values)
        self._upper: list[int] = []  # min-heap

    def add_num(self, num: int) -> None:
        """Add a number to the data structure. O(log n)."""
        # Push to lower half (negate for max-heap behavior)
        heapq.heappush(self._lower, -num)

        # Balance: lower's max must be <= upper's min
        if self._upper and (-self._lower[0]) > self._upper[0]:
            val = -heapq.heappop(self._lower)
            heapq.heappush(self._upper, val)

        # Rebalance sizes: lower can have at most 1 more element than upper
        if len(self._lower) > len(self._upper) + 1:
            val = -heapq.heappop(self._lower)
            heapq.heappush(self._upper, val)
        elif len(self._upper) > len(self._lower):
            val = heapq.heappop(self._upper)
            heapq.heappush(self._lower, -val)

    def find_median(self) -> float:
        """Return the current median. O(1)."""
        if len(self._lower) == len(self._upper):
            return (-self._lower[0] + self._upper[0]) / 2.0
        return float(-self._lower[0])  # lower has the extra element

mf = MedianFinder()
for n in [1, 2, 3, 4, 5]:
    mf.add_num(n)
    print(f"After adding {n}: median = {mf.find_median()}")
# After adding 1: median = 1.0
# After adding 2: median = 1.5
# After adding 3: median = 2.0
# After adding 4: median = 2.5
# After adding 5: median = 3.0