Heap / Priority Queue Patterns: Complete Reference¶

import heapq

# Min-heap (default)
min_heap = []
heapq.heappush(min_heap, 5)
smallest = heapq.heappop(min_heap)

# Max-heap (negate values)
max_heap = []
heapq.heappush(max_heap, -5)  # Store negative
largest = -heapq.heappop(max_heap)  # Negate on retrieval

def heap_top_k(elements: Iterable[T], k: int, key: Callable = None) -> List[T]:
    """
    Find top-k elements efficiently.

    Strategy: Maintain a min-heap of size k.
    - Elements larger than heap root replace the root
    - After processing, heap contains k largest elements

    Time: O(n log k), Space: O(k)

    Why min-heap for max problem?
    - Min-heap root is the smallest of the k largest
    - Easy to check if new element should enter the set
    - If new_element > root, it belongs in top-k; evict root
    """
    min_heap = []

    for element in elements:
        val = key(element) if key else element

        if len(min_heap) < k:
            heapq.heappush(min_heap, (val, element))
        elif val > min_heap[0][0]:
            heapq.heapreplace(min_heap, (val, element))

    return [item[1] for item in min_heap]

# Pattern: heap_kth_element
# See: docs/patterns/heap/templates.md Section 1 (Base Template)

class SolutionHeap:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        """
        Find kth largest using min-heap of size k.

        Key Insight:
        - Maintain k largest elements seen so far in a min-heap
        - Root of min-heap = smallest of the k largest = kth largest

        Why min-heap for max problem?
        - When new element > root, it belongs in top-k
        - Replace root (smallest of k largest) with new element
        - After n elements, root is exactly the kth largest
        """
        import heapq

        # Min-heap stores k largest elements
        min_heap: list[int] = []

        for num in nums:
            if len(min_heap) < k:
                # Phase 1: Fill heap until size k
                heapq.heappush(min_heap, num)
            elif num > min_heap[0]:
                # Phase 2: New element larger than kth largest
                # Replace the current kth largest
                heapq.heapreplace(min_heap, num)

        # Root is the kth largest element
        return min_heap[0]

class SolutionQuickselect:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        """
        Quickselect algorithm: O(n) average, O(n²) worst.

        Partition-based selection without full sorting.
        Use random pivot to avoid worst-case on sorted input.
        """
        import random

        target_idx = k - 1  # kth largest at index k-1 in descending order

        def partition(left: int, right: int) -> int:
            # Random pivot to avoid O(n²) on sorted arrays
            pivot_idx = random.randint(left, right)
            pivot_val = nums[pivot_idx]
            nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx]

            store_idx = left
            for i in range(left, right):
                if nums[i] >= pivot_val:  # >= for descending order
                    nums[store_idx], nums[i] = nums[i], nums[store_idx]
                    store_idx += 1

            nums[store_idx], nums[right] = nums[right], nums[store_idx]
            return store_idx

        left, right = 0, len(nums) - 1
        while left <= right:
            pivot_idx = partition(left, right)
            if pivot_idx == target_idx:
                return nums[pivot_idx]
            elif pivot_idx < target_idx:
                left = pivot_idx + 1
            else:
                right = pivot_idx - 1

        return nums[left]

# Pattern: heap_top_k
# See: docs/patterns/heap/templates.md Section 2

from collections import Counter
import heapq

class SolutionHeap:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        Find k most frequent elements using min-heap.

        Two-Phase Approach:
        1. Count frequencies: O(n)
        2. Top-k selection: O(m log k) where m = unique elements

        Why min-heap for max-frequency?
        - Store (frequency, element) tuples
        - Heap root = smallest frequency among top-k
        - Elements with higher frequency replace root
        """
        # Phase 1: Count frequencies
        frequency_map: dict[int, int] = Counter(nums)

        # Phase 2: Maintain min-heap of size k
        min_heap: list[tuple[int, int]] = []  # (frequency, element)

        for element, freq in frequency_map.items():
            if len(min_heap) < k:
                heapq.heappush(min_heap, (freq, element))
            elif freq > min_heap[0][0]:
                heapq.heapreplace(min_heap, (freq, element))

        # Extract elements (frequency is index 0, element is index 1)
        return [element for freq, element in min_heap]

class SolutionBucket:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        Bucket sort approach: O(n) time.

        Key Insight:
        - Frequency is bounded: 1 <= freq <= n
        - Use array index as frequency bucket
        - Iterate from highest frequency bucket

        Why O(n)?
        - At most n distinct elements
        - At most n buckets (frequency 1 to n)
        - Single pass through buckets
        """
        frequency_map: dict[int, int] = Counter(nums)

        # Bucket: index = frequency, value = list of elements with that frequency
        # Frequency can be at most len(nums)
        buckets: list[list[int]] = [[] for _ in range(len(nums) + 1)]

        for element, freq in frequency_map.items():
            buckets[freq].append(element)

        # Collect k elements starting from highest frequency
        result: list[int] = []
        for freq in range(len(nums), 0, -1):
            for element in buckets[freq]:
                result.append(element)
                if len(result) == k:
                    return result

        return result

# Pattern: heap_median_stream
# See: docs/patterns/heap/templates.md Section 3

import heapq

class MedianFinder:
    """
    Maintain median of a data stream using two heaps.

    Data Structure Design:
    - lower_half (max-heap): stores smaller half of numbers
    - upper_half (min-heap): stores larger half of numbers

    Invariants:
    1. All elements in lower_half <= all elements in upper_half
    2. Size difference: |lower_half| - |upper_half| <= 1
    3. If odd count, lower_half has one more element

    Why Two Heaps?
    - Median requires knowing the middle element(s)
    - Heaps give O(1) access to extreme values
    - Max-heap for lower half → O(1) access to largest of smaller numbers
    - Min-heap for upper half → O(1) access to smallest of larger numbers
    """

    def __init__(self):
        # Max-heap for lower half (store negatives for max-heap behavior)
        self.lower_half: list[int] = []  # max-heap via negation
        # Min-heap for upper half
        self.upper_half: list[int] = []  # min-heap

    def addNum(self, num: int) -> None:
        """
        Add number while maintaining heap invariants.

        Strategy:
        1. Always add to lower_half first (max-heap)
        2. Move largest from lower to upper (balances the halves)
        3. If upper becomes larger, move smallest back to lower

        This ensures:
        - Elements are correctly partitioned (lower <= upper)
        - Sizes are balanced (lower has equal or one more)
        """
        # Step 1: Add to lower half (max-heap, store negative)
        heapq.heappush(self.lower_half, -num)

        # Step 2: Move largest from lower to upper
        # This ensures lower_half max <= upper_half min
        largest_lower = -heapq.heappop(self.lower_half)
        heapq.heappush(self.upper_half, largest_lower)

        # Step 3: Rebalance if upper has more elements
        if len(self.upper_half) > len(self.lower_half):
            smallest_upper = heapq.heappop(self.upper_half)
            heapq.heappush(self.lower_half, -smallest_upper)

    def findMedian(self) -> float:
        """
        Return median in O(1) time.

        Cases:
        - Odd count: median is the root of lower_half (the extra element)
        - Even count: median is average of both roots
        """
        if len(self.lower_half) > len(self.upper_half):
            # Odd count: lower has one more
            return float(-self.lower_half[0])
        else:
            # Even count: average of two middle elements
            return (-self.lower_half[0] + self.upper_half[0]) / 2.0

from sortedcontainers import SortedList

class MedianFinderSorted:
    """
    Alternative using sorted container.
    Time: O(log n) add, O(1) find
    """

    def __init__(self):
        self.sorted_nums = SortedList()

    def addNum(self, num: int) -> None:
        self.sorted_nums.add(num)

    def findMedian(self) -> float:
        n = len(self.sorted_nums)
        mid = n // 2
        if n % 2 == 1:
            return float(self.sorted_nums[mid])
        return (self.sorted_nums[mid - 1] + self.sorted_nums[mid]) / 2.0

# Pattern: merge_k_sorted_heap
# See: docs/patterns/heap/templates.md Section 4

import heapq
from typing import List, Optional

class ListNode:
    def __init__(self, val: int = 0, next: 'ListNode' = None):
        self.val = val
        self.next = next

class SolutionHeap:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        """
        Merge k sorted lists using min-heap.

        Key Insight:
        - At any point, the next element in the merged list is the
          smallest among all current list heads
        - Min-heap gives O(log k) access to the smallest head
        - After popping, push the successor node

        Why (val, index, node) tuple?
        - val: primary sort key
        - index: tie-breaker (ListNode not comparable)
        - node: reference to actual node for result building
        """
        if not lists:
            return None

        # Min-heap: (node_value, list_index, node)
        # list_index breaks ties when values are equal
        min_heap: list[tuple[int, int, ListNode]] = []

        # Initialize heap with all non-empty list heads
        for i, node in enumerate(lists):
            if node:
                heapq.heappush(min_heap, (node.val, i, node))

        # Dummy head simplifies list construction
        dummy = ListNode(0)
        tail = dummy

        while min_heap:
            # Pop the smallest node
            val, i, node = heapq.heappop(min_heap)

            # Append to result
            tail.next = node
            tail = tail.next

            # Push successor if exists
            if node.next:
                heapq.heappush(min_heap, (node.next.val, i, node.next))

        return dummy.next

class SolutionDivideConquer:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        """
        Merge by repeatedly combining pairs.

        Analysis:
        - Each round: merge k/2 pairs → k/2 lists remain
        - After log(k) rounds: 1 list remains
        - Total: N elements × log(k) rounds = O(N log k)
        """
        if not lists:
            return None

        while len(lists) > 1:
            merged = []
            for i in range(0, len(lists), 2):
                l1 = lists[i]
                l2 = lists[i + 1] if i + 1 < len(lists) else None
                merged.append(self._merge_two(l1, l2))
            lists = merged

        return lists[0]

    def _merge_two(self, l1: Optional[ListNode], l2: Optional[ListNode]) -> Optional[ListNode]:
        dummy = ListNode(0)
        tail = dummy

        while l1 and l2:
            if l1.val <= l2.val:
                tail.next = l1
                l1 = l1.next
            else:
                tail.next = l2
                l2 = l2.next
            tail = tail.next

        tail.next = l1 if l1 else l2
        return dummy.next

# Pattern: heap_interval_scheduling
# See: docs/patterns/heap/templates.md Section 5

import heapq

class SolutionHeap:
    def minMeetingRooms(self, intervals: List[List[int]]) -> int:
        """
        Find minimum meeting rooms using min-heap of end times.

        Key Insight:
        - Sort meetings by start time (process chronologically)
        - Min-heap tracks when each room becomes free (end times)
        - For each meeting: if starts after earliest end, reuse room
        - Otherwise, allocate new room

        Why min-heap of end times?
        - We need the room that frees up earliest
        - If new meeting starts >= earliest end, room can be reused
        - Pop the old end time, push new end time
        """
        if not intervals:
            return 0

        # Sort by start time
        intervals.sort(key=lambda x: x[0])

        # Min-heap of end times (when rooms become free)
        end_times: list[int] = []

        for start, end in intervals:
            # Check if earliest-ending room is now free
            if end_times and end_times[0] <= start:
                # Room is free, reuse it (pop old end, push new end)
                heapq.heapreplace(end_times, end)
            else:
                # All rooms busy, allocate new room
                heapq.heappush(end_times, end)

        # Number of rooms = size of heap
        return len(end_times)

class SolutionSweepLine:
    def minMeetingRooms(self, intervals: List[List[int]]) -> int:
        """
        Sweep line: track events at each time point.

        Events:
        - +1 at start time (meeting begins)
        - -1 at end time (meeting ends)

        Max concurrent events = max rooms needed.
        """
        events = []
        for start, end in intervals:
            events.append((start, 1))   # Meeting starts
            events.append((end, -1))    # Meeting ends

        # Sort by time; if same time, process ends before starts
        # (room frees up before new meeting uses it)
        events.sort(key=lambda x: (x[0], x[1]))

        max_rooms = 0
        current_rooms = 0

        for time, delta in events:
            current_rooms += delta
            max_rooms = max(max_rooms, current_rooms)

        return max_rooms

# Pattern: heap_task_scheduler
# See: docs/patterns/heap/templates.md Section 6

import heapq
from collections import Counter, deque

class SolutionHeap:
    def leastInterval(self, tasks: List[str], n: int) -> int:
        """
        Schedule tasks with cooldown using max-heap.

        Strategy:
        1. Greedily pick most frequent available task
        2. Track cooldown using a queue (task becomes available at time t)
        3. If no task available, idle (time still advances)

        Why max-heap?
        - We want to process high-frequency tasks first
        - This minimizes idle time at the end
        - Max-heap gives O(log k) access to most frequent
        """
        # Count task frequencies
        task_counts = Counter(tasks)

        # Max-heap of remaining counts (negate for max-heap)
        max_heap = [-count for count in task_counts.values()]
        heapq.heapify(max_heap)

        # Queue of (available_time, remaining_count) for tasks in cooldown
        cooldown_queue: deque[tuple[int, int]] = deque()

        time = 0

        while max_heap or cooldown_queue:
            time += 1

            # Check if any task exits cooldown
            if cooldown_queue and cooldown_queue[0][0] == time:
                available_time, remaining = cooldown_queue.popleft()
                heapq.heappush(max_heap, -remaining)

            if max_heap:
                # Execute most frequent available task
                count = -heapq.heappop(max_heap)
                count -= 1

                if count > 0:
                    # Task has more instances, put in cooldown
                    cooldown_queue.append((time + n + 1, count))
            # else: idle (no task available)

        return time

class SolutionMath:
    def leastInterval(self, tasks: List[str], n: int) -> int:
        """
        Mathematical approach based on most frequent task.

        Observation:
        - Most frequent task determines the "frame"
        - Frame: chunks of size (n+1) with one instance of max-freq task

        Formula:
        - frames = (max_freq - 1)
        - frame_size = n + 1
        - base_time = frames * frame_size
        - Add tasks with max_freq (they fill the last partial frame)

        Edge case:
        - If total tasks > formula result, no idle needed
        """
        task_counts = Counter(tasks)
        max_freq = max(task_counts.values())
        max_freq_count = sum(1 for count in task_counts.values() if count == max_freq)

        # Formula: (max_freq - 1) frames × (n + 1) slots + max_freq_count final slots
        formula_time = (max_freq - 1) * (n + 1) + max_freq_count

        # Actual time is max of formula and total tasks
        return max(formula_time, len(tasks))

# Pattern: heap_greedy_simulation
# See: docs/patterns/heap/templates.md Section 7

import heapq

class SolutionHeap:
    def lastStoneWeight(self, stones: List[int]) -> int:
        """
        Simulate stone smashing using max-heap.

        Game Rules:
        1. Pick two heaviest stones x, y (x <= y)
        2. If x == y: both destroyed
        3. If x != y: stone of weight y - x remains

        Why max-heap?
        - Need repeated access to two largest elements
        - Python heapq is min-heap, so negate values
        - O(log n) per operation vs O(n) for linear search
        """
        # Convert to max-heap (negate all values)
        max_heap = [-stone for stone in stones]
        heapq.heapify(max_heap)  # O(n) heapify

        # Simulate until 0 or 1 stone remains
        while len(max_heap) > 1:
            # Pop two largest (remember to negate)
            largest = -heapq.heappop(max_heap)
            second = -heapq.heappop(max_heap)

            # If different weights, push the difference
            if largest != second:
                heapq.heappush(max_heap, -(largest - second))

        # Return last stone weight (or 0 if no stones left)
        return -max_heap[0] if max_heap else 0

import heapq

def find_kth_largest(nums: List[int], k: int) -> int:
    """Find kth largest element. Time: O(n log k), Space: O(k)"""
    min_heap = []
    for num in nums:
        if len(min_heap) < k:
            heapq.heappush(min_heap, num)
        elif num > min_heap[0]:
            heapq.heapreplace(min_heap, num)
    return min_heap[0]

from collections import Counter
import heapq

def top_k_frequent(nums: List[int], k: int) -> List[int]:
    """Find k most frequent elements. Time: O(n + m log k)"""
    freq = Counter(nums)
    min_heap = []
    for element, count in freq.items():
        if len(min_heap) < k:
            heapq.heappush(min_heap, (count, element))
        elif count > min_heap[0][0]:
            heapq.heapreplace(min_heap, (count, element))
    return [elem for _, elem in min_heap]

import heapq

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

    def add(self, num: int) -> None:
        heapq.heappush(self.lower, -num)
        heapq.heappush(self.upper, -heapq.heappop(self.lower))
        if 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.0

import heapq

def merge_k_sorted(lists: List[List[int]]) -> List[int]:
    """Merge k sorted lists. Time: O(N log k)"""
    min_heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(min_heap, (lst[0], i, 0))

    result = []
    while min_heap:
        val, list_idx, elem_idx = heapq.heappop(min_heap)
        result.append(val)
        if elem_idx + 1 < len(lists[list_idx]):
            next_val = lists[list_idx][elem_idx + 1]
            heapq.heappush(min_heap, (next_val, list_idx, elem_idx + 1))
    return result

import heapq

def min_rooms(intervals: List[List[int]]) -> int:
    """Minimum rooms for non-overlapping meetings. Time: O(n log n)"""
    if not intervals:
        return 0
    intervals.sort(key=lambda x: x[0])
    end_times = []
    for start, end in intervals:
        if end_times and end_times[0] <= start:
            heapq.heapreplace(end_times, end)
        else:
            heapq.heappush(end_times, end)
    return len(end_times)

import heapq

def simulate_max_heap(items: List[int]) -> int:
    """Process items by repeatedly taking largest. Time: O(n log n)"""
    max_heap = [-x for x in items]
    heapq.heapify(max_heap)

    while len(max_heap) > 1:
        largest = -heapq.heappop(max_heap)
        second = -heapq.heappop(max_heap)
        # Process largest and second
        result = process(largest, second)
        if result > 0:
            heapq.heappush(max_heap, -result)

    return -max_heap[0] if max_heap else 0

import heapq

# Create heap
heap = []                    # Empty heap
heapq.heapify(list)          # Convert list to heap O(n)

# Add element
heapq.heappush(heap, val)    # Add element O(log n)

# Remove smallest
smallest = heapq.heappop(heap)  # Remove and return O(log n)

# Peek smallest
smallest = heap[0]           # View without removing O(1)

# Push then pop (optimized)
heapq.heappushpop(heap, val) # Push val, pop smallest O(log n)

# Pop then push (optimized)
heapq.heapreplace(heap, val) # Pop smallest, push val O(log n)

# Max-heap using negation
max_heap = [-x for x in items]
heapq.heapify(max_heap)
largest = -heapq.heappop(max_heap)