Priority Queue Patterns — Top-K, Merge K Lists, Median, Dijkstra

A priority queue (implemented as a heap) gives O(log n) insert and O(1) access to the min or max element. Python’s heapq module provides a min-heap. For a max-heap, negate the values.

import heapq

# Min-heap basics
h = []
heapq.heappush(h, 3)
heapq.heappush(h, 1)
heapq.heappush(h, 2)
heapq.heappop(h)      # → 1 (smallest)

# Max-heap trick: negate values
heapq.heappush(h, -5)
-heapq.heappop(h)     # → 5 (largest)

Pattern 1: Top-K Elements

Find the K largest (or smallest) elements in a collection.

Key insight: Use a min-heap of size K to find the K largest elements. The heap’s minimum is the K-th largest.

K-th Largest Element (LeetCode 215)

import heapq

def find_kth_largest(nums: list[int], k: int) -> int:
    """
    Find k-th largest element.
    Time: O(n log k), Space: O(k)
    """
    heap = []

    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)  # remove smallest

    return heap[0]  # k-th largest is the smallest in our k-sized heap

Top-K Frequent Elements (LeetCode 347)

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 log k), Space: O(n)
    """
    freq = Counter(nums)
    # Min-heap of size k, keyed by frequency
    return heapq.nlargest(k, freq.keys(), key=freq.get)

Manual approach (more interview-friendly):

def top_k_frequent_manual(nums: list[int], k: int) -> list[int]:
    freq = Counter(nums)
    heap = []

    for num, count in freq.items():
        heapq.heappush(heap, (count, num))
        if len(heap) > k:
            heapq.heappop(heap)

    return [num for count, num in heap]

Why Min-Heap for Largest?

Goal	Heap Type	Size	Pop When
K largest	Min-heap	K	Smallest leaves, K largest stay
K smallest	Max-heap	K	Largest leaves, K smallest stay

Pattern 2: Merge K Sorted Lists

Merge K sorted lists into one sorted list. This appears as LeetCode 23.

import heapq

def merge_k_sorted(lists: list[list[int]]) -> list[int]:
    """
    Merge K sorted lists.
    Time: O(N log K) where N = total elements
    Space: O(K) for the heap
    """
    heap = []
    result = []

    # Initialize with first element from each list
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))
            # (value, list_index, element_index)

    while heap:
        val, list_idx, elem_idx = heapq.heappop(heap)
        result.append(val)

        # Push next element from the same list
        if elem_idx + 1 < len(lists[list_idx]):
            next_val = lists[list_idx][elem_idx + 1]
            heapq.heappush(heap, (next_val, list_idx, elem_idx + 1))

    return result

Trace: Merge [[1,4,7], [2,5,8], [3,6,9]]

Heap init: [(1,0,0), (2,1,0), (3,2,0)]

Pop (1,0,0) → result=[1], push (4,0,1)
  Heap: [(2,1,0), (3,2,0), (4,0,1)]

Pop (2,1,0) → result=[1,2], push (5,1,1)
  Heap: [(3,2,0), (4,0,1), (5,1,1)]

Pop (3,2,0) → result=[1,2,3], push (6,2,1)
  Heap: [(4,0,1), (5,1,1), (6,2,1)]

... continues until all merged
Result: [1,2,3,4,5,6,7,8,9] ✓

For Linked Lists (LeetCode 23)

import heapq

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

def merge_k_lists(lists):
    """
    Merge K sorted linked lists.
    """
    heap = []
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))

    dummy = tail = ListNode(0)

    while heap:
        val, idx, node = heapq.heappop(heap)
        tail.next = node
        tail = tail.next
        if node.next:
            heapq.heappush(heap, (node.next.val, idx, node.next))

    return dummy.next

Pattern 3: Running Median (Two Heaps)

Maintain the median of a stream of numbers using a max-heap for the lower half and a min-heap for the upper half. This is LeetCode 295.

import heapq

class MedianFinder:
    """
    Find median from data stream — LeetCode 295.
    Time: O(log n) per add, O(1) per find
    Space: O(n)
    """
    def __init__(self):
        self.low = []   # max-heap (negated) — smaller half
        self.high = []  # min-heap — larger half

    def add_num(self, num: int) -> None:
        # Always add to max-heap first
        heapq.heappush(self.low, -num)

        # Ensure max of low <= min of high
        if self.low and self.high and -self.low[0] > self.high[0]:
            val = -heapq.heappop(self.low)
            heapq.heappush(self.high, val)

        # Balance sizes (low can have at most 1 extra)
        if len(self.low) > len(self.high) + 1:
            val = -heapq.heappop(self.low)
            heapq.heappush(self.high, val)
        elif len(self.high) > len(self.low):
            val = heapq.heappop(self.high)
            heapq.heappush(self.low, -val)

    def find_median(self) -> float:
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

Trace

add(1):  low=[-1]     high=[]        median=1
add(2):  low=[-1]     high=[2]       median=(1+2)/2=1.5
add(3):  low=[-2,-1]  high=[3]
         rebalance → low=[-2] high=[3]
         wait — add 3: push -3 to low → low=[-3,-1], high=[2]
         -low[0]=3 > high[0]=2 → move: low=[-2,-1] high=[2,3]
         len check: low=2, high=2 → balanced
         median = (2+2)/2 = 2.0

Correct trace:
add(1): low=[-1]         high=[]        median=1.0
add(2): push -2 to low → low=[-2,-1]
        len(low)=2 > len(high)+1=1 → move 2 to high
        low=[-1]         high=[2]       median=1.5
add(3): push -3 to low → low=[-3,-1]
        -low[0]=3 > high[0]=2 → move: low=[-2,-1] high=[2,3]
        but now low=2, high=2 → balanced
        wait that's wrong too. Let me be precise:

add(3): push -3 to low → low=[-3,-1], high=[2]
        -low[0]=3 > high[0]=2 → pop 3 from low, push to high
        low=[-1], high=[2,3]
        len(high)=2 > len(low)=1 → pop 2 from high, push to low
        low=[-2,-1], high=[3]
        median = (2+3)/2 = 2.5? No...

Actually median of [1,2,3] = 2. Let me re-examine.

The key invariants:

1. max(low) <= min(high) — smaller half and larger half
2. |len(low) - len(high)| <= 1 — balanced sizes
3. If odd count, low has the extra element

Simpler Trace

Numbers: [5, 2, 8, 1, 9]

add(5): low=[-5]        high=[]         → median=5
add(2): low=[-2]        high=[5]        → median=3.5
add(8): low=[-5,-2]     high=[8]        → median=5
add(1): low=[-2,-1]     high=[5,8]      → median=3.5
add(9): low=[-5,-2,-1]  high=[8,9]      → median=5

Pattern 4: Dijkstra’s Shortest Path

Use a min-heap as a priority queue to always process the closest unvisited node.

import heapq

def dijkstra(graph: dict, start: int) -> dict:
    """
    Single-source shortest paths.
    Time: O((V+E) log V), Space: O(V)
    """
    dist = {start: 0}
    heap = [(0, start)]  # (distance, node)

    while heap:
        d, u = heapq.heappop(heap)

        # Skip if we already found a shorter path
        if d > dist.get(u, float('inf')):
            continue

        for v, weight in graph.get(u, []):
            new_dist = d + weight
            if new_dist < dist.get(v, float('inf')):
                dist[v] = new_dist
                heapq.heappush(heap, (new_dist, v))

    return dist

Trace

Graph: 0→1(4), 0→2(1), 2→1(2), 1→3(1), 2→3(5)

Heap: [(0,0)]   dist={0:0}

Pop (0,0): process neighbors
  → 1: dist=4   heap=[(1,2),(4,1)]  dist={0:0, 1:4, 2:1}
  → 2: dist=1

Pop (1,2): process neighbors
  → 1: 1+2=3 < 4 → update!  dist={0:0, 1:3, 2:1}
  → 3: 1+5=6                dist={0:0, 1:3, 2:1, 3:6}

Pop (3,1): process neighbors
  → 3: 3+1=4 < 6 → update!  dist={0:0, 1:3, 2:1, 3:4}

Pop (4,1): d=4 > dist[1]=3 → skip (stale entry)

Pop (4,3): d=4 = dist[3] → process (no better neighbors)

Pop (6,3): d=6 > dist[3]=4 → skip

Result: {0:0, 1:3, 2:1, 3:4} ✓

Pattern Summary

Pattern	Heap Type	Size	Use Case
Top-K	Min-heap	K	K largest elements, K-th largest
Merge K Lists	Min-heap	K	Merge sorted sequences
Running Median	Two heaps	n	Stream median, sliding window median
Dijkstra	Min-heap	up to E	Shortest paths, weighted graphs

Edge Cases

# Top-K with k = n (return all)
# Top-K with k = 1 (just find max)
# Merge K with empty lists
# Median with single element
# Dijkstra with disconnected graph

When to Use This Pattern

• “K largest/smallest” → Top-K with opposite heap
• “Merge sorted” → Heap selects next minimum
• “Median of stream” → Two heaps split data
• “Shortest path” → Dijkstra with priority queue
• “Schedule by priority” → Direct priority queue use

Related Problems

• Heap Advanced Patterns — more heap techniques
• Sliding Window Maximum — can also use heap
• Deque Design Patterns — alternative to heap for 0-1 BFS