题目
给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
定义一对值 (u,v),其中第一个元素来自 nums1,第二个元素来自 nums2 。
请找到和最小的 k 个数对 (u1,v1), (u2,v2) ... (uk,vk) 。
示例 1:
输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
示例 2:
输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
示例 3:
输入: nums1 = [1,2], nums2 = [3], k = 3
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
提示:
1 <= nums1.length, nums2.length <= 10^5
-10^9 <= nums1[i], nums2[i] <= 10^9
nums1 和 nums2 均为升序排列
1 <= k <= 10^4
优先队列思路
一开始想的太简单了,错误思路:不能把俩个数组里的全部数的组合都放小根堆里去,因为俩个数组的最大容量是10^ 5,那么所有组合的最大数量为10^10 肯定会爆内存。
那么如何优化? 关键是题目已经给出俩个数组都是升序的。首先nums1[0]和nums1[0]肯定是最小的组合,次小的组合肯定是(nums1[1],nums1[0])或者(nums1[0],nums1[1])。因为俩个数组都是升序,所以(nums1[1],nums1[1])>=(nums1[1],nums1[0])、(nums1[0],nums1[1])
所以我们可以把前面较小的一部分组合的下标放入小根堆,然后根据下标对应的值的大小来排序,如果(i,j)下标对应的值此时是最小的,就把其取出存入答案链表,并把次小(i+1,j)、(i,j+1)下标放入小根堆。
但是这样有问题,就是(i,j+1)和(i+1,j)被取出时都会加入(i+1,j+1),导致这个坐标被重复加入了,所以我们可以改变思路,当(i,j)此时是最小的,就把其取出,只把次小(i,j+1)放入小根堆,至于(i+1,j) 则等当取出(i+1,j-1)时再加入。
这种取法要求我们一开始就往小根堆加入(0,0),(1,0)...,(m,0),这样才能保证所有的点位能够被加入。
用图解来看就是,从下面这张图的策略转为下下张图的策略。
代码:
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
PriorityQueue<int[]>priorityQueue=new PriorityQueue<>((o1,o2)->
nums1[o1[0]]+nums2[o1[1]]-nums1[o2[0]]-nums2[o2[1]]);//自定义排序小根堆 lambda写法
int m=nums1.length,n=nums2.length;
for (int i = 0; i < Math.min(m,k); i++) {//Math.min(m,k)是优化写法 如果k>=m则不用说 第一列全加入就好了
priorityQueue.add(new int[]{nums1[i],nums2[0]});//当k<m 只加到第k行就行了 这是为什么呢?
}//因为只要求前k小 前k行包括的可能性肯定足够了 再往下的行都比前面的大 不可能是答案
List<List<Integer>> list=new ArrayList<>();
for (int i = 0; i < k; i++) {
int []o=priorityQueue.poll();
list.add(new ArrayList<>(Arrays.asList(nums1[o[0]],nums2[o[1]])));
if(o[1]+1<n) priorityQueue.add(new int[]{o[0],o[1]+1});
if(priorityQueue.isEmpty()) break;
}
return list;
}