[Algorithm] Find K Pairs with Smallest Sums

Find K Pairs with Smallest Sums

Problem statement:

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.

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

Solution 1: Naive approach.

Time complexity: O(nmlog(k))

Space complexity: O(k)

The naive approach is to get all possible pairs and use a heap with max size of K to find the top K smallest, which will perform in O(nmlog(k)). The implementation is quite straight forward, we just need to iterate through the two input arrays, get combination of the two and push into the max heap, when the heap hits size K+1, pop out the maximum item.

This approach works fine with unsorted arrays and given that the two input arrays are sorted, we have this condition that we are not utilizing, so there must be a faster way.

Solution 2: N way merging + Heap.

Time complexity: O(klog(n))

Space complexity: O(n)

Since the two arrays are sorted we have the following condition:

n[i] + m[j] < n[i] + m[j+1] for all i in [0 ... n-1]

Let's use this notation for sum of n[i] and m[j]:

s[i][j] = n[i] + m[j]

And -i to denote any possible index in n except i and j* to denote any possible j in m.

And when collecting top K sums, we need to compare s[i][j] with s[-i][j*], we already know that if i is fixed for n, then n[i][j*] has an ascending order. 

If we view the sequence of s[i][j*] as a linked list, then the comparison becomes the comparison between heads of each s[i*]. Thus we can use a min heap to hold s[i*] and pop out the smallest from the heap and replace it with s[i][j+1] which is the next element in the sequence of s[i], very much similar to a N-way merging. The key here is how we advance the sequence.

Code implementation:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42

class Solution { public: vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) { std::priority_queue<Record> heap; int n = nums1.size(); if (n == 0) return {}; int m = nums2.size(); if (m == 0) return {}; for (int i=0; i<n; i++) { heap.push(Record(nums1[i] + nums2[0], nums1[i])); } vector<pair<int, int>> results; while (results.size() < k && !heap.empty()) { Record top = heap.top(); heap.pop(); results.push_back(pair<int, int>(top.ref, top.value-top.ref)); if (top.pos+1 < m) { top.pos++; top.value = top.ref + nums2[top.pos]; heap.push(top); } } return results; } struct Record { int value; int ref; int pos; Record(int v, int r) : value(v), ref(r), pos(0) {} bool operator<(const Record& rhs) const { return value > rhs.value; } }; };