← Back to all patterns

K-way Merge and Two Heaps

A heap answers one question fast: what is the smallest thing right now? K-way merge asks it of K sorted lists at once. Two heaps ask it from both sides of the middle, which gives you a running median.

Pattern 9 used a heap to keep the K best items. It named two other heap moves and left them for later. This page covers both. They share one tool, heapq, and one habit: push tuples with a safe tie-break so Python never has to compare two objects it cannot order.

When to use

• K sorted inputs. Lists, files, rows of a matrix, or log streams. Each one is sorted, but they are not sorted together. Use K-way merge.
• The k-th item of a merged order. Merge lazily and stop after k pops. You never build the full merge.
• A running middle. “Add a number, then report the median”, many times. Use two heaps.
• Pick from two pools by a rule. IPO and Sliding Window Median also split items into two heaps and move items between them.

Core idea

K-way merge: a heap of heads

Reading the figure. Each row is one sorted list. Dashed cells are already in the output. The amber cells are the current heads, and the heap holds exactly those three as (value, list). The heap top is the next item to emit. After it leaves, only list 2 moves forward, so only one push follows.

Two heaps: one for each half

Reading the figure. The left tree is the lower half, with its largest value on top. The right tree is the upper half, with its smallest value on top. The amber tops face each other across the dashed median line. Every value on the left is at most every value on the right. So the median comes from the tops alone.

The templates

import heapq


def merge_sorted_arrays(arrays: list[list[int]]) -> list[int]:
    """All values from K sorted lists, in one sorted list.

    Example:
        >>> merge_sorted_arrays([[1, 4, 7], [2, 5], [3, 6, 9]])
        [1, 2, 3, 4, 5, 6, 7, 9]
    """
    # One entry per non-empty list: (head value, list index, position).
    # The list index is the tie-break. Position 0 is each list's head.
    heap = [(arr[0], index, 0) for index, arr in enumerate(arrays) if arr]
    heapq.heapify(heap)           # O(K) to build, cheaper than K pushes

    merged: list[int] = []

    # Invariant: heap holds the next unused item of every list not yet empty,
    # so heap[0] is the smallest unused item overall.
    while heap:
        value, index, pos = heapq.heappop(heap)
        merged.append(value)

        # pos + 1 is the next item in the same list. Push it if it exists.
        if pos + 1 < len(arrays[index]):
            heapq.heappush(heap, (arrays[index][pos + 1], index, pos + 1))

    return merged

In production code, list(heapq.merge(*arrays)) does the same thing lazily. Write the loop in the interview, then name heapq.merge.

def running_medians(stream: list[int]) -> list[float]:
    """The median after each new number arrives.

    Example:
        >>> running_medians([5, 15, 1, 3])
        [5.0, 10.0, 5.0, 4.0]
    """
    lower: list[int] = []         # max-heap of the small half, stored negated
    upper: list[int] = []         # min-heap of the large half
    medians: list[float] = []

    # Invariant at the top of each pass: every value in lower is at most every
    # value in upper, and len(lower) is len(upper) or len(upper) + 1.
    for value in stream:
        # Step 1: push into lower (negate it, max-heap). Then move lower's
        # largest across, so upper gets the right value and order holds.
        heapq.heappush(lower, -value)
        heapq.heappush(upper, -heapq.heappop(lower))   # - undoes the negation

        # Step 2: rebalance. Upper may now be one too big. Move its smallest
        # back, negated, so lower is never smaller than upper.
        if len(upper) > len(lower):
            heapq.heappush(lower, -heapq.heappop(upper))

        # Step 3: read the tops. lower[0] is the negated max of the low half.
        if len(lower) > len(upper):
            medians.append(float(-lower[0]))           # odd count: lower top
        else:
            # Even count: average the two tops. / 2 is true division.
            medians.append((-lower[0] + upper[0]) / 2)

    return medians

Reading the figure. Each column is one moment while 1 is added. Blue cells are the lower half and violet cells the upper half. The amber cell is the value that just moved. The move up always sends the largest value of the lower half, so the order between halves holds. The move back only fixes the sizes.

The push-then-move dance in step 1 is the trick. It avoids comparing the new value with either top by hand, and it can never break the order between halves.

Common mistakes

• Pushing (value, node). On a tie, Python compares the nodes and raises TypeError. Put an int index in the middle.
• Pushing every item up front. That is a heap of N items and O(N log N). The point is a heap of K heads, so O(N log K).
• Forgetting to negate on the way out of the lower max-heap. lower[0] is the negated max. Read it as -lower[0].
• Rebalancing by size only. If a new number goes straight into the smaller heap, the order between halves can break. Route it through the other heap first.
• Integer division for the median. // 2 on 1 + 2 gives 1, not 1.5. Use / 2.
• Skipping empty lists. arr[0] on an empty list raises IndexError. Filter them out when you seed the heap.

The problems

from __future__ import annotations

import heapq


class MergeNode:
    """A singly linked list node for this page's merge problems."""

    def __init__(self, val: int = 0, next: MergeNode | None = None) -> None:
        self.val = val            # default 0: only the dummy head uses it
        self.next = next


def build_merge_list(values: list[int]) -> MergeNode | None:
    """Linked list holding values, in order. Empty input gives None."""
    head = None
    # Build from the back, so each new node points at the one built before.
    for value in reversed(values):
        head = MergeNode(value, head)
    return head


def merge_list_values(head: MergeNode | None) -> list[int]:
    """The values of a linked list, as a Python list."""
    out: list[int] = []
    while head is not None:
        out.append(head.val)
        head = head.next
    return out


def merge_k_lists(lists: list[MergeNode | None]) -> MergeNode | None:
    """Merge k sorted linked lists into one sorted linked list.

    Args:
        lists: Heads of sorted linked lists. Any of them may be None.

    Returns:
        The head of the merged list, or None if every list is empty.

    Example:
        >>> heads = [build_merge_list([1, 4, 5]), build_merge_list([1, 3, 4]),
        ...          build_merge_list([2, 6])]
        >>> merge_list_values(merge_k_lists(heads))
        [1, 1, 2, 3, 4, 4, 5, 6]
        >>> merge_k_lists([None, None]) is None
        True
    """
    # (value, list index, node). The index breaks ties, so heapq never
    # has to compare two MergeNode objects, which would raise TypeError.
    heap: list[tuple[int, int, MergeNode]] = []
    for index, head in enumerate(lists):
        if head is not None:      # skip empty lists: they have no head
            heap.append((head.val, index, head))
    heapq.heapify(heap)           # O(k) build

    dummy = MergeNode()           # placeholder before the real head
    tail = dummy                  # last node of the merged list so far

    # Invariant: dummy.next .. tail is sorted and holds every popped node.
    # heap holds the current head of every list not yet used up.
    while heap:
        _, index, node = heapq.heappop(heap)   # _ : value, already in node
        tail.next = node
        tail = node

        # Same index: the next node comes from the same list, and that list
        # has no other entry in the heap, so the tie-break stays unique.
        if node.next is not None:
            heapq.heappush(heap, (node.next.val, index, node.next))

    return dummy.next             # skip the placeholder

def kth_smallest_in_matrix(matrix: list[list[int]], k: int) -> int:
    """The k-th smallest value in a matrix with sorted rows and columns.

    Args:
        matrix: A non-empty n x n matrix. Rows and columns ascend.
        k: Rank to find, 1 <= k <= n * n.

    Returns:
        The k-th smallest value, duplicates counted.

    Example:
        >>> kth_smallest_in_matrix([[1, 5, 9], [10, 11, 13], [12, 13, 15]], 8)
        13
        >>> kth_smallest_in_matrix([[-5]], 1)
        -5
    """
    n = len(matrix)

    # Seed with column 0 of each row: (value, row, col). Rows past k cannot
    # hold the answer, so min(n, k) rows is enough. row is the tie-break.
    heap = [(matrix[row][0], row, 0) for row in range(min(n, k))]
    heapq.heapify(heap)

    # Pop k - 1 times: after that the smallest k - 1 are gone and the top
    # of the heap is the k-th smallest.
    # Invariant: heap holds the next unused cell of every seeded row.
    for _ in range(k - 1):
        _, row, col = heapq.heappop(heap)
        # col + 1 is the next cell to the right. Push it if the row has one.
        if col + 1 < n:
            heapq.heappush(heap, (matrix[row][col + 1], row, col + 1))

    return heap[0][0]             # heap[0]: smallest entry. [0]: its value

class MedianFinder:
    """Running median with two heaps.

    lower is a max-heap (values negated) holding the smaller half.
    upper is a min-heap holding the larger half. len(lower) is always
    len(upper) or len(upper) + 1, and every lower value <= every upper value.

    Example:
        >>> finder = MedianFinder()
        >>> finder.add_num(1)
        >>> finder.add_num(2)
        >>> finder.find_median()
        1.5
        >>> finder.add_num(3)
        >>> finder.find_median()
        2.0
    """

    def __init__(self) -> None:
        self.lower: list[int] = []    # negated values: lower[0] is -(max of low half)
        self.upper: list[int] = []    # plain values: upper[0] is min of high half

    def add_num(self, num: int) -> None:
        """Add num to the stream in O(log n)."""
        # Push into lower, negated so heapq's min acts as a max.
        heapq.heappush(self.lower, -num)
        # Move lower's largest to upper. - undoes the negation. This keeps
        # every lower value <= every upper value, whatever num was.
        heapq.heappush(self.upper, -heapq.heappop(self.lower))

        # Upper may now hold one more than lower. Move its smallest back,
        # negated, so lower is equal or exactly 1 bigger.
        if len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def find_median(self) -> float:
        """Median of every number added so far, in O(1).

        Raises:
            ValueError: If no number has been added.
        """
        if not self.lower:            # lower is empty only when both are
            raise ValueError("no numbers added yet")

        # Odd count: lower holds the extra one, and its top is the median.
        # [0] is the heap top, and - turns the stored negative back.
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])

        # Even count: average the two middle values. / 2 is true division,
        # so 1 and 2 give 1.5, not 1.
        return (-self.lower[0] + self.upper[0]) / 2

def smallest_range(nums: list[list[int]]) -> list[int]:
    """Smallest [lo, hi] holding at least one value from every list.

    Args:
        nums: k non-empty lists, each sorted ascending.

    Returns:
        [lo, hi] for the smallest such range. Ties go to the smaller lo.

    Raises:
        ValueError: If nums is empty or any list is empty.

    Example:
        >>> smallest_range([[4, 10, 15, 24, 26], [0, 9, 12, 20], [5, 18, 22, 30]])
        [20, 24]
        >>> smallest_range([[1, 2, 3], [1, 2, 3]])
        [1, 1]
    """
    if not nums or any(not row for row in nums):
        raise ValueError("need at least one list, and no empty lists")

    # One pick per list: (value, list index, position). Position 0 is the
    # list's smallest value. The list index is the tie-break.
    heap = [(row[0], index, 0) for index, row in enumerate(nums)]
    heapq.heapify(heap)
    # The heap gives the min. The max has to be tracked by hand.
    current_max = max(row[0] for row in nums)       # [0]: each list's head

    # Start with the infinite range, so the first real range always wins.
    best_lo, best_hi = float("-inf"), float("inf")

    # Invariant: heap holds exactly one value from every list, and
    # current_max is the largest of them. So [heap[0][0], current_max]
    # covers every list.
    while True:
        low, index, pos = heapq.heappop(heap)

        # Strict <: on a tie keep the earlier range, whose lo is smaller,
        # because lo only grows as the loop runs.
        if current_max - low < best_hi - best_lo:
            best_lo, best_hi = low, current_max

        # pos + 1 == len: this list has no next value, so no later pick can
        # cover it. Every remaining range has been tried.
        if pos + 1 == len(nums[index]):
            return [best_lo, best_hi]

        nxt = nums[index][pos + 1]                  # pos + 1: next in this list
        heapq.heappush(heap, (nxt, index, pos + 1))
        current_max = max(current_max, nxt)         # the new pick may be the max

Recap

Where this goes next

Pattern 21, Weighted Shortest Paths, puts the same heap of (cost, tie-break, node) tuples at the centre of Dijkstra’s algorithm. There the heap holds the frontier of a graph instead of the heads of K lists.