Problem description:

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u, v) which consists of one element from the first array and one element from the second array.

Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.

Example 1:

1
2
3
Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

1
2
3
Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

1
2
3
Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Constraints:

  • 1 <= nums1.length, nums2.length <= 104
  • 109 <= nums1[i], nums2[i] <= 109
  • nums1 and nums2 both are sorted in ascending order.
  • 1 <= k <= 1000

Solution:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# # # # # ? . . 
# # # ? . . . .
# ? . . . . . .   "#" means pair already in the output
# ? . . . . . .   "?" means pair currently in the queue
# ? . . . . . .
? . . . . . . .
. . . . . . . .

nums1: [-1, 2, 3, 4]
nums2: [-5, 1, 2, 5]

   -1  2    3    4
-5  5  -10  -15  -20
1  -1  -2   ...
2  ... 
5  ...

Think like a 2d array, we want to use a priority queue to store some elements to find current minimum.

If we pop out element with j = 0, it means nums[j] would create smaller element, so we keep add column[0] element to heap 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        queue = []
        def push(i, j):
            if i < len(nums1) and j < len(nums2):
                heapq.heappush(queue, [nums1[i] + nums2[j], i, j])
        push(0, 0)
        res = []
        while queue and len(res) < k:
            print("------")
            _, i, j = heapq.heappop(queue)
            print("cur: ", i, j)
            res.append([nums1[i], nums2[j]])
            push(i, j + 1)
            print(i, j+1)
            if j == 0:
                push(i + 1, 0)
                print(i+1, 0)
        return res

time complexity: $O(klogk)$
space complexity: $O()$
reference:
related problem: