373. Find K Pairs with Smallest Sums

Summary

You are given two integer arrays nums1 and nums2, sorted in non-decreasing order, and an integer k.

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

Return the k pairs with the smallest sums.

Approach

Use a heap to keep the smallest candidates.

Each item in the heap looks like:

(sum, u, v)

The min-heap always keeps the current smallest candidate.

The inputs are already sorted, so we can start from (0, 0).

Then we add the next candidates: (u + 1, v) and (u, v + 1).

This requires only k iterations to build the answer.

We also need to track visited index pairs because the same pair can be reached more than once:

start: (0, 0)
next: (1, 0), (0, 1)

Those two nodes can both reach (1, 1), so we need to skip it if it was already visited.

Function

Initialization:

heap=[(nums1[0]+nums2[0], 0,0)]
visited = {(0, 0)}
ans = []

While len(ans) < k: Pop (sum, u, v) from the heap root. Add (nums1[u], nums2[v]) to ans. For (nu, nv) in [(u + 1, v), (u, v + 1)]: If nu < len(nums1) and nv < len(nums2) and (nu, nv) is not visited: Push (nums1[nu] + nums2[nv], nu, nv) into the heap. Mark (nu, nv) as visited.

Return ans.

Complexity

Time Complexity: O(K log K) Space Complexity: O(K)