Pattern 12: K-Way Merge (Min & Max Heap / Priority Queue)

The K-Way Merge pattern is a fundamental algorithmic technique designed to efficiently combine $K$ sorted data streams (such as arrays, linked lists, or matrices) into a single, globally sorted structure. By leveraging a Min-Heap (Priority Queue) to track the smallest unmerged element across all streams, this pattern reduces the time complexity from a brute-force $O(N \cdot K)$ or standard sorting $O(N \log N)$ down to an optimal $O(N \log K)$, where $N$ is the total number of elements across all $K$ inputs.

1. Understanding the K-Way Merge Pattern

Motivation & Problem Context

When merging two sorted arrays (as in the classic Merge Sort divide phase), a simple two-pointer approach suffices: compare the current heads of both arrays and pick the smaller one in $O(1)$ time.

However, when scaled to $K$ sorted arrays or linked lists, naively checking the current heads of all $K$ lists at every step requires scanning $K$ elements per step. Across $N$ total elements, this produces a naive time complexity of $O(N \cdot K)$.

The Min-Heap Solution

Instead of linearly scanning all $K$ heads, we insert the head of each list into a Min-Heap of max size $K$.

1. Finding the absolute minimum element across all $K$ lists takes $O(1)$ time (the root of the Min-Heap).
2. Extracting the minimum element and pushing its successor from the same list into the heap takes $O(\log K)$ time.

Thus, processing all $N$ elements takes $O(N \log K)$ time, offering massive scalability when $K$ grows large.

2. Core Mental Model & State Tracking

To correctly implement a K-Way Merge, the Min-Heap must store more than just raw element values. When the smallest value is popped from the heap, the algorithm must know which array or list that element came from so it can advance that specific list's pointer to fetch the next candidate.

Tuple Representation in Heap

Depending on the data structure used, every entry pushed into the Min-Heap carries metadata:

• For Arrays: A triplet or node structure containing (value, array_index, element_index).
• For Linked Lists: The pointer directly to the ListNode (since node.next preserves the reference to the successor node).
• For Pair Sums / Matrix Traversal: A triplet containing (sum/value, row_index / i, col_index / j).

3. General Algorithmic Step-by-Step Workflow

The standard execution pipeline for any K-Way Merge problem follows three distinct phases:

Phase 1: Bounded Initialization

1. Initialize a Min-Heap configured with a custom comparator based on element values.
2. Iterate through all $K$ input lists and insert their first (smallest) element into the Min-Heap.
3. If $K$ is larger than the total number of lists or requested items, bound the initialization to $\min(K, \text{total lists})$.

Phase 2: Heap-Driven Extraction Loop

While the Min-Heap is not empty:

1. Extract: Pop the root element (the absolute minimum currently available across all active streams).
2. Append: Append the extracted element's value to the output list/array.
3. Advance: Check if the list/array from which the extracted element originated contains a subsequent element.
   • If node.next != null (or index + 1 < length), push the next element into the Min-Heap.
   • If no further elements exist in that stream, do nothing; the active heap size naturally decreases.

Phase 3: Termination

The loop completes automatically when all input streams are exhausted and the heap becomes empty.

4. Problem 1: Merge K Sorted Lists (LeetCode 23)

Problem Statement

Given an array of $K$ linked-lists lists, where each linked-list is sorted in ascending order, merge all the linked-lists into one sorted linked-list and return it.

Step-by-Step Visual Trace

Given $K = 3$ sorted linked lists:

• List 0: $1 \to 4 \to 5$
• List 1: $1 \to 3 \to 4$
• List 2: $2 \to 6$

5. Implementation & Complexity: Merge K Sorted Lists

Complete Java Implementation

import java.util.PriorityQueue;

/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode() {}
 *     ListNode(int val) { this.val = val; }
 *     ListNode(int val, ListNode next) { this.val = val; this.next = next; }
 * }
 */
public class MergeKSortedLists {

    public ListNode mergeKLists(ListNode[] lists) {
        // Base Condition: Handle null or empty input array
        if (lists == null || lists.length == 0) {
            return null;
        }

        // Initialize Min-Heap with a custom comparator comparing node values
        PriorityQueue<ListNode> minHeap = new PriorityQueue<>(
            (a, b) -> Integer.compare(a.val, b.val)
        );

        // Phase 1: Push the head node of every non-null list into the Min-Heap
        for (ListNode head : lists) {
            if (head != null) {
                minHeap.offer(head);
            }
        }

        // Dummy head pattern simplifies result list construction
        ListNode dummyHead = new ListNode(0);
        ListNode current = dummyHead;

        // Phase 2: Extract minimum node and push its successor
        while (!minHeap.isEmpty()) {
            // Extract the smallest node currently in the heap
            ListNode minNode = minHeap.poll();

            // Attach minNode to the merged result list
            current.next = minNode;
            current = current.next;

            // If a successor exists in the same list, push it to the heap
            if (minNode.next != null) {
                minHeap.offer(minNode.next);
            }
        }

        // Return the head of the newly merged sorted list
        return dummyHead.next;
    }
}

Mathematical Complexity Analysis

Note: $N$ represents the total number of nodes across all $K$ lists ($N = \sum_{i=0}^{K-1} \text{length}(lists[i])$).

6. Problem 2: Find K Pairs with Smallest Sums (LeetCode 373)

Problem Statement

Given two integer arrays nums1 and nums2 sorted in non-descending order and an integer k, define a pair $(u, v)$ consisting of one element $u \in \text{nums1}$ and one element $v \in \text{nums2}$. Return the $k$ pairs $(u, v)$ with the smallest sums.

Mathematical & Conceptual Formulation

Consider matrix $M$ where $M[i][j] = \text{nums1}[i] + \text{nums2}[j]$.

Because both nums1 and nums2 are sorted, every row $i$ and every column $j$ in this implicit matrix is sorted in non-descending order. Finding the $K$ smallest pairs is equivalent to performing a K-Way Merge across sorted rows of this matrix.

nums1 = [1, 7, 11], nums2 = [2, 4, 6]

Implicit Sum Matrix M[i][j]:
            nums2[0]=2   nums2[1]=4   nums2[2]=6
nums1[0]=1    (1+2)=3      (1+4)=5      (1+6)=7
nums1[1]=7    (7+2)=9      (7+4)=11     (7+6)=13
nums1[2]=11   (11+2)=13    (11+4)=15    (11+6)=17

The Bounded Heap Optimization

To avoid generating all $O(\vert nums1\vert \cdot \vert nums2\vert )$ pairs:

1. Seed the Min-Heap with only the first element of each row: (nums1[i] + nums2[0], i, 0).
2. Limit the seed size to $\min(\vert nums1\vert , K)$ elements. (If we only need $K$ pairs, examining beyond the $K$-th element of nums1 is redundant).
3. When pair (i, j) is popped from the heap, the immediate successor candidate in that row is (i, j + 1).

7. Deep-Dive Code Evolution: Find K Pairs with Smallest Sums

Complete Java Implementation

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;

public class FindKPairsWithSmallestSums {

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();

        // Base Conditions
        if (nums1 == null || nums2 == null || nums1.length == 0 || nums2.length == 0 || k <= 0) {
            return result;
        }

        // Min-Heap stores array of 3 integers: [sum, index_in_nums1, index_in_nums2]
        // Comparator sorts elements by sum (index 0) ascending
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(
            (a, b) -> Integer.compare(a[0], b[0])
        );

        // Optimization: Bound initial seed size to min(nums1.length, k)
        int initialRows = Math.min(nums1.length, k);

        // Seed Phase: Push (nums1[i] + nums2[0], i, 0) into heap
        for (int i = 0; i < initialRows; i++) {
            minHeap.offer(new int[] { nums1[i] + nums2[0], i, 0 });
        }

        // Extraction Phase: Extract at most k pairs
        while (k > 0 && !minHeap.isEmpty()) {
            int[] current = minHeap.poll();
            int i = current[1]; // Index in nums1
            int j = current[2]; // Index in nums2

            // Add extracted pair [nums1[i], nums2[j]] to output list
            result.add(Arrays.asList(nums1[i], nums2[j]));
            k--; // Decrement required pairs count

            // Successor Generation: Check if next column (j + 1) exists in nums2
            if (j + 1 < nums2.length) {
                int nextSum = nums1[i] + nums2[j + 1];
                minHeap.offer(new int[] { nextSum, i, j + 1 });
            }
        }

        return result;
    }
}

Step-by-Step Dry Run

Inputs: nums1 = [1, 7, 11], nums2 = [2, 4, 6], k = 3.

1. Initialization

• initialRows = Math.min(3, 3) = 3.
• Offer nums1[0] + nums2[0] = $1+2=3 \implies [3, 0, 0]$.
• Offer nums1[1] + nums2[0] = $7+2=9 \implies [9, 1, 0]$.
• Offer nums1[2] + nums2[0] = $11+2=13 \implies [13, 2, 0]$.
• minHeap state: [[3, 0, 0], [9, 1, 0], [13, 2, 0]].

2. Iteration 1 ($k = 3$)

• poll() $\implies [3, 0, 0]$ ($i=0, j=0$).
• Append [nums1[0], nums2[0]] = [1, 2] to result. Decrement $k \to 2$.
• Successor check: $j+1 = 1 < \text{nums2.length}$ (4 exists).
• Push [nums1[0] + nums2[1], 0, 1] = $[1+4, 0, 1] = [5, 0, 1]$.
• minHeap state: [[5, 0, 1], [9, 1, 0], [13, 2, 0]].

3. Iteration 2 ($k = 2$)

• poll() $\implies [5, 0, 1]$ ($i=0, j=1$).
• Append [nums1[0], nums2[1]] = [1, 4] to result. Decrement $k \to 1$.
• Successor check: $j+1 = 2 < \text{nums2.length}$ (6 exists).
• Push [nums1[0] + nums2[2], 0, 2] = $[1+6, 0, 2] = [7, 0, 2]$.
• minHeap state: [[7, 0, 2], [9, 1, 0], [13, 2, 0]].

4. Iteration 3 ($k = 1$)

• poll() $\implies [7, 0, 2]$ ($i=0, j=2$).
• Append [nums1[0], nums2[2]] = [1, 6] to result. Decrement $k \to 0$.
• Loop terminates ($k = 0$).

Final Output: [[1, 2], [1, 4], [1, 6]].

Complexity Matrix for K Smallest Pairs

8. Tabular Synthesis of K-Way Merge Variations

9. Common Edge Cases & Production Debugging Pitfalls

1. Arithmetic Integer Overflow

• Issue: When computing pair sums (e.g., $u + v$), if elements in nums1 and nums2 are large integers near Integer.MAX_VALUE, $u + v$ can overflow and produce negative values, breaking custom comparator logic.
• Fix: Use Integer.compare(a[0], b[0]) or cast sum computations to long. Avoid subtraction logic like (a, b) -> a.val - b.val if values can span Integer.MIN_VALUE to Integer.MAX_VALUE.

// SAFE Comparator
PriorityQueue<int[]> minHeap = new PriorityQueue<>(
    (a, b) -> Integer.compare(a[0], b[0])
);

// UNSAFE Comparator (Can overflow if a[0] is negative or b[0] is negative)
PriorityQueue<int[]> minHeap = new PriorityQueue<>(
    (a, b) -> a[0] - b[0]
);

2. Null or Empty Sub-lists

• Issue: In ListNode[] lists, individual entries can be null (e.g., [[], [1, 2], []]).
• Fix: Always guard the initial seeding loop with an explicit non-null check (if (head != null)).

3. Pointer Out of Bounds

• Issue: In array-based K-Way Merge problems, advancing the index pointer without checking array bounds leads to ArrayIndexOutOfBoundsException.
• Fix: Ensure j + 1 < array.length before pushing successor elements into the Min-Heap.

4. Bounding Heap Initialization to $K$

• Issue: Seeding the Min-Heap with $1,000,000$ array heads when $K = 5$ wastes $\mathcal{O}(M)$ space and computation.
• Fix: Always use Math.min(array.length, k) during the initial seeding pass.

10. Summary and Key Takeaways

1. Core Purpose: K-Way Merge optimizes combining $K$ pre-sorted streams down to a logarithmic time complexity factor relative to $K$ ($\mathcal{O}(N \log K)$).
2. Min-Heap Utility: The Min-Heap acts as an active window storing only the smallest unconsumed element from each active stream.
3. State Tracking: Heap entries must preserve source metadata (e.g., list pointers or array indices i, j) to identify which stream must advance upon extraction.
4. List Merging Efficiency: When merging linked lists, mutating pointer links directly (current.next = minNode) eliminates secondary memory allocations, keeping auxiliary space bounded to $\mathcal{O}(K)$.
5. Matrix & Pair Sum Abstraction: Problems asking for $K$ smallest pairs or sums across sorted arrays map directly to K-Way Merge by treating each row as a sorted stream.
6. Bounded Seeding Optimization: For pair problems where $K$ is smaller than input lengths, initialize the heap with only $\min(\text{length}, K)$ elements to save time and space.
7. Comparator Safety: Always prefer standard numeric comparators like Integer.compare() to prevent silent integer overflow errors in priority queues.