[LeetCode 373] 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.
Find 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]
Key observation: K pairs && smallest. -> Heap
Min Heap solution. O(K * log K) runtime, O(K) space
1. create min heap that takes a pair's index and compares pairs' sum.
2. add the first Math.min(nums1.length, k) pairs in format (i, 0) to min heap. i is nums1's index, 0 represents 0 index of nums2.
3. poll from min heap, add the pair to answer, then create a new pair with nums2 index incremented by 1 and add it back to min heap.
4. repeat 3 until either there are k pairs or there is no pairs left in min heap.
A variation of this algorithm is to only add the smallest pair(0, 0) to min heap. Then each time we poll a pair, we try to increase the pair's 1st index; then try to increase the pair's 2nd index. However, we'll need some extra book keeping to avoid adding the same pair more than once.
class Solution { public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) { List<List<Integer>> res = new ArrayList<>(); if(nums1.length == 0 || nums2.length == 0) { return res; } PriorityQueue<int[]> minPq = new PriorityQueue<>((a, b) -> { return nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]]; }); Set<Integer>[] set = new Set[nums1.length]; for(int i = 0; i < nums1.length; i++) { set[i] = new HashSet<>(); } minPq.add(new int[]{0, 0}); set[0].add(0); while(k > 0 && minPq.size() > 0) { int[] curr = minPq.poll(); List<Integer> list = new ArrayList<>(); list.add(nums1[curr[0]]); list.add(nums2[curr[1]]); res.add(list); if(curr[0] < nums1.length - 1 && !set[curr[0] + 1].contains(curr[1])) { minPq.add(new int[]{curr[0] + 1, curr[1]}); set[curr[0] + 1].add(curr[1]); } if(curr[1] < nums2.length - 1 && !set[curr[0]].contains(curr[1] + 1)) { minPq.add(new int[]{curr[0], curr[1] + 1}); set[curr[0]].add(curr[1] + 1); } k--; } return res; } }
Related Problems
[LeetCode 373] Find K Pairs with Smallest Sums
[LeetCode 719] Find K-th Smallest Pair Distance
[LeetCode 1439] Find the Kth Smallest Sum of a Matrix With Sorted Rows