Find K Pairs with Smallest Sums

---

Find K Pairs with Smallest Sums: 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:

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:

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:

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 <= 10^5
• -10^9 <= nums1[i], nums2[i] <= 10^9
• nums1 and nums2 both are sorted in ascending order.
• 1 <= k <= 10^4

---

Try this Problem on your own or check similar problems:

1. Kth Smallest Element in a Sorted Matrix
2. Kth Smallest Product of Two Sorted Arrays
3. Find K-th Smallest Pair Distance

Solution:

Java

public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<List<Integer>> result = new ArrayList<>();
    PriorityQueue<int[]> p = new PriorityQueue<>((a, b) -> a[2] - b[2]);
    int m = nums1.length, n = nums2.length;

    for(int i = 0; i < Math.min(k, n); ++i){
      p.add(new int[]{0, i, nums1[0] + nums2[i]});
    }

    for(int i = 0;  i < Math.min(k, (long) n*m); ++i){
      int[] current = p.poll();
      int nums1Index = current[0], nums2Index = current[1];
      result.add(List.of(nums1[nums1Index], nums2[nums2Index]));
      if(nums1Index == m - 1) continue;
      p.add(new int[]{nums1Index + 1, nums2Index, nums1[nums1Index + 1] + nums2[nums2Index]});
    }

    return result;
}

JavaScript

/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number} k
 * @return {number[][]}
 */
var kSmallestPairs = function (nums1, nums2, k) {
  const minHeap = new MinPriorityQueue({ priority: (x) => x[0] });
  const n = nums2.length;
  const result = [];

  for (let i = 0; i < nums1.length; i++) {
    const num1 = nums1[i];
    const num2 = nums2[0];

    minHeap.enqueue([num1 + num2, i, 0]);
  }

  while (k > 0 && !minHeap.isEmpty()) {
    const [sum, idx1, idx2] = minHeap.dequeue().element;

    result.push([nums1[idx1], nums2[idx2]]);

    if (result.length === k) return result;

    if (idx2 < n - 1) {
      minHeap.enqueue([nums1[idx1] + nums2[idx2 + 1], idx1, idx2 + 1]);
    }
  }

  return result;
};

Python

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        result = []
        pq = []
        m = len(nums1)
        n = len(nums2)

        for i in range(min(k, m)):
            heapq.heappush(pq, [nums1[i] + nums2[0], i, 0])

        while pq and len(result) < k:
            _, nums1Index, nums2Index = heapq.heappop(pq)
            result.append([nums1[nums1Index], nums2[nums2Index]])
            if nums2Index + 1 < n:
                heapq.heappush(pq, [nums1[nums1Index] + nums2[nums2Index + 1], nums1Index, nums2Index + 1])

        return result

C++

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> result;
        int n = nums1.size();
        int m = nums2.size();
        priority_queue<tuple<int, int, int>, vector<tuple<int, int, int >>, greater<tuple<int, int, int>>> pq;
        for(int i = 0; i < n; i++) {
            pq.push({nums1[i] + nums2[0], i, 0});
        }
        while(k-- and pq.size()) {
            auto [sum, i, j] = pq.top(); pq.pop();
            result.push_back({nums1[i], nums2[j]});
            if(j < m - 1)
                pq.push({nums1[i] + nums2[j + 1], i, j + 1});
        }
        return result;
    }
};

---

Time/Space Complexity:

• Time Complexity: O(klogk)
• Space Complexity: O(k)

---

Explanation:

We can reuse the same algorithm as described in Kth Smallest Element in a Sorted Matrix by choosing the first set of candidates to be the pairs generated by combining the smallest number from the first array with all elements in the second, than we can extract the minimum pair k times (each time storing it in the result). When we poll the last pair containing first element in first array the next best thing is to choose the next element in the first array (nums1Index + 1). Note that we're not exactly extracting the top element from the queue k times but min(k, m*n) times since the k can be greater than the number of pairs we can generate (we use (long) to cope with the overflow if m and n are large numbers). Time complexity is O(klogk) since we iterate over k times (you could put a stricter boundary O(min(k, m*n))) and we extract an element from a queue (operations on the queue are log k where k is the size of the queue) with size k (or to be stricter min(k,n)). The space complexity is the size of the queue.