Find K Pairs with Smallest Sums — Heap on an Implicit Sorted Matrix

Problem Statement#

Given two sorted arrays nums1 and nums2 and an integer k, return the k pairs (a, b) with the smallest sums where a is from nums1 and b is from nums2.

Constraints:

• 1 <= nums1.length, nums2.length <= 10^5
• -10^9 <= nums1[i], nums2[i] <= 10^9
• 1 <= k <= 10^4

Input:  nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]

Input:  nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]

Why This Problem Matters#

Find K Pairs with Smallest Sums is a Google and Amazon priority queue interview favorite. The naive O(nm log(nm)) approach blows up with arrays of 10^5 elements. The optimal heap solution runs in O(k log k) and is conceptually identical to k-way merge on a sorted matrix.

The pattern shows up everywhere — k-th smallest in a sorted matrix, k smallest products, and pair selection in dual-sorted streams.

The Core Insight#

Imagine the implicit matrix M[i][j] = nums1[i] + nums2[j]. Both rows and columns are sorted. Use the same min-heap pattern as LeetCode 378 — push the first column (one entry per row of nums1, each paired with j = 0), then on every pop push the right neighbor.

Cap the initial seed to the first k rows; we never need more.

Visual Dry Run#

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

Solution (Optimal)#

import heapq
 
class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        if not nums1 or not nums2:
            return []
        heap = []
        for i in range(min(k, len(nums1))):
            heapq.heappush(heap, (nums1[i] + nums2[0], i, 0))
        res = []
        while heap and len(res) < k:
            s, i, j = heapq.heappop(heap)
            res.append([nums1[i], nums2[j]])
            if j + 1 < len(nums2):
                heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
        return res

var kSmallestPairs = function(nums1, nums2, k) {
    if (!nums1.length || !nums2.length) return [];
    const heap = new MinHeap((a, b) => a[0] - b[0]);
    for (let i = 0; i < Math.min(k, nums1.length); i++) {
        heap.push([nums1[i] + nums2[0], i, 0]);
    }
    const res = [];
    while (heap.size() && res.length < k) {
        const [, i, j] = heap.pop();
        res.push([nums1[i], nums2[j]]);
        if (j + 1 < nums2.length) heap.push([nums1[i] + nums2[j + 1], i, j + 1]);
    }
    return res;
};

Time: O(k log k) — heap holds at most k items, k pops. Space: O(k) — heap and result.

Common Mistakes#

• Generating all n*m sums and sorting — TLE on large inputs
• Seeding the heap with all of nums1 even when k is small — wastes time
• Forgetting to handle empty input arrays
• Pushing the down-neighbor in addition to the right-neighbor — duplicates entries
• Mishandling duplicate sums — they are valid; tuple tiebreakers prevent crashes

Interview Tips#

• Draw the implicit sum matrix on the whiteboard — interviewers love this
• Discuss why pushing only the right neighbor is enough (because we seed the entire first column)
• Mention the symmetric variant: seed first row instead, push down neighbor
• Add a guard against duplicate cells using a visited set when generalizing to other problems

Follow-up Questions#

• What if you must push both right and down? Add a visited set to dedupe (i, j)
• K smallest products instead of sums? Same heap with the multiplication
• Three sorted arrays for triples? Same idea, tuple becomes (sum, i, j, k)
• Streaming nums2? Use a sorted data structure for incremental insertions
• What if k is much larger than nm? Cap k at nm up front

Key Takeaways#

• Implicit sorted matrix view turns pair sums into a k-way merge problem
• Seed first column with min(k, n) entries — don't over-fill the heap
• Push only the right neighbor when seeding the column (or only down when seeding the row)
• O(k log k) time, O(k) space — beats O(nm log(nm)) brute force
• Same template solves k-th smallest in sorted matrix and k-th smallest product
• Tuples with index tiebreakers prevent comparison errors
• Heap FAANG pattern — memorize both seeding directions