Problem 373

Posted by Ruizhi Ma on October 24, 2020

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;

    }
}