Solution URL
https://leetcode.com/submissions/detail/412764304/
代码
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
//ans: https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/solution/javada-ding-dui-xiao-ding-dui-jie-fa-yi-dong-by-vi/
//建立大顶堆
PriorityQueue<List<Integer>> pq = new PriorityQueue<>(k, (o1, o2) -> {
return (o2.get(1) + o2.get(0)) - (o1.get(1) + o1.get(0));
});
//两层for循环将数字加入堆中
for(int i = 0; i < Math.min(k, nums1.length); i++){
for(int j = 0; j < Math.min(k, nums2.length); j++){
//剪枝,因为是nums都是顺序排列的,所以当堆已经满了,并且数组之和大于堆顶元素,则此后所有都不符合条件
if(k == pq.size() && nums1[i] + nums2[j] > pq.peek().get(0) + pq.peek().get(1)){
break;
}
//此时过来的都符合条件,应该加入堆中
//如果已经满了,则弹出一个
if(k == pq.size()) pq.poll();
//加入堆中
List<Integer> pair = new ArrayList<>();
pair.add(nums1[i]);
pair.add(nums2[j]);
pq.add(pair);
}
}
//遍历堆,倒序加入结果集(因为大顶堆)
List<List<Integer>> res = new ArrayList<>();
for(int i = 0; i < k && !pq.isEmpty(); i++){//这里一定要加!pq.isEmpty(),防止k给的过大;否则会有空节点被弹出
res.add(0, pq.poll());
}
return res;
}
}