Find K Pairs with Smallest Sums

LeetCode 373: Find K Pairs with Smallest Sums Solution & Explanation

Problem Overview

TL;DR: The optimal solution treats the problem as merging $N$ sorted lists using a Min-Heap to efficiently select the $k$ smallest sums without generating all pairs.

The problem asks us to find the $k$ pairs with the smallest sums formed by picking one integer from array nums1 and one from array nums2. Both input arrays are sorted in non-decreasing order. This is a classic optimization problem found in interviews, often referred to as LeetCode 373. The goal is to return these pairs efficiently, avoiding the computational cost of processing every possible combination.

Brute Force Approach for Find K Pairs with Smallest Sums

The most straightforward approach is to generate every possible pair $(u, v)$ where $u \in \text{nums1}$ and $v \in \text{nums2}$, calculate their sums, and find the smallest $k$.

Naive Algorithm Steps

1. Initialize an empty list all_pairs.
2. Iterate through every element $i$ in nums1.
3. For each $i$, iterate through every element $j$ in nums2.
4. Append the pair [nums1[i], nums2[j]] to all_pairs.
5. Sort all_pairs based on the sum of elements.
6. Return the first $k$ elements.

Pseudo-code

pairs = []
for u in nums1:
    for v in nums2:
        pairs.append([u, v])
sort pairs by (u + v)
return pairs[0 to k]

Analysis of Failure

• Time Complexity: $O(N \cdot M \log(N \cdot M))$, where $N$ and $M$ are the lengths of nums1 and nums2. Since constraints allow lengths up to $10^5$, $N \cdot M$ can reach $10^{10}$. Sorting this many elements is computationally infeasible.
• Space Complexity: $O(N \cdot M)$ to store all pairs. This will immediately cause a Memory Limit Exceeded error on large inputs.
• Why it fails: This approach ignores the critical property that the input arrays are already sorted. It performs redundant work by calculating large sums that have no chance of being in the top $k$.

Algorithm Strategy: Heap (Priority Queue) Patterns

We utilize a Min-Heap to maintain the "frontier" of candidates.

1. Initialization: Instead of adding all $N \times M$ pairs, we only initialize the heap with the first pair from each "row". Specifically, we push (nums1[i] + nums2[0], i, 0) for the first $k$ elements of nums1. We only need the first $k$ rows because the element at nums1[k] (combined with nums2[0]) is guaranteed to be larger than any of the first $k$ pairs if nums1 is sorted.
2. Tracking State: The heap stores tuples containing (sum, index_in_nums1, index_in_nums2). The sum determines the priority.
3. Iterative Extraction:
   • Extract the element with the smallest sum from the heap. Let's say this corresponds to indices $(i, j)$.
   • Add this pair to our result list.
   • Expand: If there is a next element in that specific row (i.e., nums2 has an element at $j+1$), calculate the new sum nums1[i] + nums2[j+1] and push (new_sum, i, j+1) back into the heap.
4. Termination: Repeat until we have collected $k$ pairs or the heap is empty.

This strategy enforces that we only generate and compare candidates that are plausibly the next smallest, keeping the heap size small.

Execution Flow

1. Validation: Check if nums1 or nums2 is empty or if $k=0$. Return empty list if so.
2. Heap Setup: Create a Min-Heap. Iterate $i$ from $0$ to $\min(k, \text{length of nums1})$. Push (nums1[i] + nums2[0], i, 0) onto the heap.
3. Selection Loop: While the result list size is less than $k$ and the heap is not empty:
   • Pop the top element: (current_sum, i, j).
   • Append [nums1[i], nums2[j]] to the result list.
   • Check if $j + 1 < \text{length of nums2}$.
   • If valid, push (nums1[i] + nums2[j+1], i, j+1) into the heap.
4. Return: Return the accumulated result list.

Proof of Correctness

The correctness relies on the invariant that the heap always contains the smallest unvisited element from every active row.

• Initially, the heap contains the minimum element of every relevant row (the first column).
• The global minimum must be the root of the Min-Heap.
• When we remove (nums1[i], nums2[j]), the next smallest element in row $i$ is (nums1[i], nums2[j+1]).
• Since nums2 is sorted, nums1[i] + nums2[j+1] \ge nums1[i] + nums2[j].
• By adding this successor to the heap, we ensure the heap continues to hold the "head" of every sorted stream. Therefore, successive pops yield the globally smallest sums in order.

Pattern Reuse Notes

The Heap - K-way Merge pattern used in LeetCode 373 is highly versatile. Understanding this solution helps solve:

• LeetCode 23: Merge k Sorted Lists - The canonical version of this pattern where the input is explicit linked lists rather than implicit array combinations.
• LeetCode 378: Kth Smallest Element in a Sorted Matrix - Identical logic. The rows of the matrix are the sorted lists.
• LeetCode 632: Smallest Range Covering Elements from K Lists - Uses the same heap structure to maintain a window across $k$ sorted lists.

Reference Implementation

C++ Solution for LeetCode 373

1#include <vector>
2#include <queue>
3#include <tuple>
4
5using namespace std;
6
7class Solution {
8public:
9    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
10        vector<vector<int>> result;
11        if (nums1.empty() || nums2.empty() || k == 0) {
12            return result;
13        }
14
15        // Min-heap stores tuples: {sum, index_in_nums1, index_in_nums2}
16        // Utilizing greater<> for min-heap behavior
17        auto comp = [&](const tuple<int, int, int>& a, const tuple<int, int, int>& b) {
18            return get<0>(a) > get<0>(b);
19        };
20        priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, decltype(comp)> minHeap(comp);
21
22        // Initialize heap with the first element of the first k rows
23        // We only need to consider the first k elements of nums1
24        for (int i = 0; i < nums1.size() && i < k; ++i) {
25            minHeap.push({nums1[i] + nums2[0], i, 0});
26        }
27
28        while (k-- > 0 && !minHeap.empty()) {
29            auto [sum, i, j] = minHeap.top();
30            minHeap.pop();
31
32            result.push_back({nums1[i], nums2[j]});
33
34            // If there is a next element in the current row (nums1[i]), push it
35            if (j + 1 < nums2.size()) {
36                minHeap.push({nums1[i] + nums2[j + 1], i, j + 1});
37            }
38        }
39
40        return result;
41    }
42};

Java Solution for LeetCode 373

1import java.util.ArrayList;
2import java.util.List;
3import java.util.PriorityQueue;
4
5class Solution {
6    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
7        List<List<Integer>> result = new ArrayList<>();
8        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
9            return result;
10        }
11
12        // Min-heap stores int array: [sum, index_nums1, index_nums2]
13        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
14
15        // Initialize heap with the first element of the first k rows from nums1
16        // combined with nums2[0]
17        for (int i = 0; i < nums1.length && i < k; i++) {
18            minHeap.offer(new int[]{nums1[i] + nums2[0], i, 0});
19        }
20
21        while (k > 0 && !minHeap.isEmpty()) {
22            int[] current = minHeap.poll();
23            int i = current[1];
24            int j = current[2];
25
26            List<Integer> pair = new ArrayList<>();
27            pair.add(nums1[i]);
28            pair.add(nums2[j]);
29            result.add(pair);
30
31            // Push the next element from the same row (increment index in nums2)
32            if (j + 1 < nums2.length) {
33                minHeap.offer(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
34            }
35            
36            k--;
37        }
38
39        return result;
40    }
41}

Python Solution for LeetCode 373

1import heapq
2
3class Solution:
4    def kSmallestPairs(self, nums1: list[int], nums2: list[int], k: int) -> list[list[int]]:
5        if not nums1 or not nums2 or k == 0:
6            return []
7        
8        # Min-heap stores tuples: (sum, index_nums1, index_nums2)
9        min_heap = []
10        
11        # Initialize heap with the first element of the first k rows
12        # We pair nums1[i] with nums2[0]
13        for i in range(min(k, len(nums1))):
14            heapq.heappush(min_heap, (nums1[i] + nums2[0], i, 0))
15            
16        result = []
17        
18        while k > 0 and min_heap:
19            current_sum, i, j = heapq.heappop(min_heap)
20            result.append([nums1[i], nums2[j]])
21            
22            # If there is a next element in the current row (nums1[i] + nums2[j+1])
23            if j + 1 < len(nums2):
24                heapq.heappush(min_heap, (nums1[i] + nums2[j + 1], i, j + 1))
25                
26            k -= 1
27            
28        return result