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]
改进：参考 https://leetcode.com/problems/find-k-pairs-with-smallest-sums/discuss/84551/simple-Java-O(KlogK)-solution-with-explanation
因为两个数组都是有序的，往min heap中添加元素的理想顺序应该是：首先只考虑两个数组的前k个元素，对应nums1中的某个元素，下一个理想顺序的组合应该是 nums1中的这个元素 + 从堆中poll出来的对应的nums2元素的下一个元素。堆中存index比较方便，先把(i, 0)放入堆中，i对应nums1的index，0是nums2的index；poll出来一个最小值之后，保存结果，检查index是否越界，然后把下一个candidate放入堆中
basic idea: use minHeap to keep track on next minimum pair sum, and we only need to maintain K possible candidates in the data structure.
for every numbers in nums1, its best candidate(yields min sum) always starts from nums2[0] since arrays are all sorted; and for a specific number in nums1, its next candidate sould be [this specific number] + nums2[current_associated_index + 1], unless out of boundary
时间：O(KlogK)，空间：O(K)
class Solution { public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) { List<int[]> res = new ArrayList<>(); if(nums1.length == 0 || nums2.length == 0 || k == 0) return res; PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]])); // for(int i = 0; i < nums1.length && i < k; i++) { // minHeap.offer(new int[] {i, 0}); // } // while(k-- > 0 && !minHeap.isEmpty()) { // int[] tmp = minHeap.poll(); // res.add(new int[] {nums1[tmp[0]], nums2[tmp[1]]}); // if(tmp[1] == nums2.length - 1) // continue; // minHeap.offer(new int[] {tmp[0], tmp[1] + 1}); // } // return res; for(int i = 0; i < nums2.length && i < k; i++) { minHeap.offer(new int[] {0, i}); } while(k-- > 0 && !minHeap.isEmpty()) { int[] tmp = minHeap.poll(); res.add(new int[] {nums1[tmp[0]], nums2[tmp[1]]}); if(tmp[0] == nums1.length - 1) continue; minHeap.offer(new int[] {tmp[0]+ 1, tmp[1]}); } return res; } }
之前的方法用了min heap，但类似于brute force，能过但是太慢了，O(K^2 logK)
class Solution { public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) { List<int[]> res = new ArrayList<>(); if(nums1 == null || nums2 == null) return res; PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (b[0] + b[1]) - (a[0] + a[1])); int m = Math.min(nums1.length, k), n = Math.min(nums2.length, k); for(int i = 0; i < m; i++) { for(int j = 0; j < n; j++) { minHeap.offer(new int[] {nums1[i], nums2[j]}); if(minHeap.size() > k) minHeap.poll(); } } while(!minHeap.isEmpty()) { res.add(0, minHeap.poll()); } return res; } }
or
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 || k == 0) { return res; } PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]])); boolean[][] visited = new boolean[nums1.length][nums2.length]; minHeap.offer(new int[] {0, 0}); visited[0][0] = true; for(int i = 0; i < k && !minHeap.isEmpty(); i++) { int[] cur = minHeap.poll(); int r = cur[0], c = cur[1]; List<Integer> tmp = new ArrayList<>(); tmp.add(nums1[r]); tmp.add(nums2[c]); res.add(tmp); if(r + 1 < nums1.length && !visited[r + 1][c]) { minHeap.offer(new int[] {r + 1, c}); visited[r + 1][c] = true; } if(c + 1 < nums2.length && !visited[r][c + 1]) { minHeap.offer(new int[] {r, c + 1}); visited[r][c + 1] = true; } } return res; } }