K Pairs with Smallest Sums

Problem

Given two sorted lists and k, return the k pairs [u, v] with the smallest sums, ordered by sum and then by index.

Seed a min-heap with every (nums1[i], nums2[0]) pair and, after popping pair (i, j), push (i, j + 1), which visits each row in increasing order. Return fewer than k pairs when they are exhausted.

Hint

For each row of nums1 the sums increase along nums2, so one heap per row is enough.

Test cases

case 1

in: [[1,7,11],[2,4,6],3]

expected: [[1,2],[1,4],[1,6]]

case 2

in: [[1,1,2],[1,2,3],2]

expected: [[1,1],[1,1]]

case 3

in: [[1,2],[3],3]

expected: [[1,3],[2,3]]

case 4

in: [[],[],3]

expected: []

case 5

in: [[1],[2],5]

expected: [[1,2]]

Starter codePython

def k_smallest_pairs(nums1, nums2, k):
    # Your code here
    pass

5 test cases · run your code against them in the editor