Find K Pairs with Smallest Sums

Patrick Leaton

Published: June 6, 2021, 10:15 a.m.

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:

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

The description was taken from https://leetcode.com/problems/find-k-pairs-with-smallest-sums/,

Problem Solution

#O(N(Log K)) Time, O(K) Space

class Solution:

def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:

if not nums1 or not nums2:

return []

n, m = len(nums1), len(nums2)

seen, heap, output = {}, [], []

heapq.heappush(heap, (nums1[0] + nums2[0], 0, 0))

while heap and k > 0:

small_sum, ind_one, ind_two = heapq.heappop(heap)

output.append([nums1[ind_one], nums2[ind_two]])

k -= 1

if ind_one + 1 < n and (ind_one + 1, ind_two) not in seen:

heappush(heap, ((nums1[ind_one + 1] + nums2[ind_two]), ind_one+1, ind_two))

seen[((ind_one + 1, ind_two))] = True

if ind_two + 1 < m and (ind_one, ind_two + 1) not in seen:

heappush(heap, ((nums1[ind_one] + nums2[ind_two+1]), ind_one, ind_two+1))

seen[((ind_one, ind_two + 1))] = True

return output

Problem Explanation

Right off the bat, the title has "find k smallest" in it. That is typically an indication that a heap can be used to perform this task because it implies that we will be working off some kind of sorting.

That is what we're going to do.

Within a heap node, we will keep the index of the two values from both lists that we are summing as well as their sum, which is what the heap will sort on.

We will continuously pop the smallest sum off of the heap and append the two numbers used to create the sum into the output. We will then sum the current element of each opposing array with the next element of each current array and if the combination hasn't been summed before, we will push those values into the heap. Once have done this k times, we are done.

First off, if one of the two arrays is empty then we can't sum together one number from each array, so we will return an empty array.

if not nums1 or not nums2:

return []

Next, we will create an n and an m value to indicate the length of each array, since we don't really want to be making that calculation in each iteration.

n, m = len(nums1), len(nums2)

Afterward, we will initialize our seen hashmap, our heap, and our output array.

seen, heap, output = {}, [], []

Then, we will push the sum of the first two elements from the input arrays into the heap along with their indices.

Since the arrays are already sorted in ascending order, we would expect these two to be the first smallest sum.

heapq.heappush(heap, (nums1[0] + nums2[0], 0, 0))

Now, we will utilize a loop that will run while there are elements in the heap and while we haven't found k sums.

while heap and k > 0:

At the beginning of each iteration, we will pop the smallest sum and two index values of the elements that created this sum off of the heap.

small_sum, ind_one, ind_two = heapq.heappop(heap)

Then, we will append the two elements of this small sum to the output by using the index values we just popped.

output.append([nums1[ind_one], nums2[ind_two]])

Afterward, we will decrement k since we just found another small sum.

k -= 1

Next, we will push two more tuples into the array.

The smartest move is to test the next element in each array, along with the current element from the other array.

If we have these two inputs, [1,2,3,4], and [5,6,7,8], with k being four, we would be making all sum combinations with one first before moving to two, [1,5], [1,6], [2,5], and [1,7]. We can use the same number multiple times as long as the combination hasn't been used before, that may not be clear from the description.

That is the approach we will want to use when pushing the smallest possible sum values into the heap.

For the first array, if the index value after the current element isn't past the last index of the array, and it hasn't already been combined with the current element of the second array, then we will push the sum of these two elements and their indices into the heap.