Find K Pairs with Smallest Sums 查找和最小的K对数字

Find K Pairs with Smallest Sums 查找和最小的K对数字

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.

Return the k pairs (u1, v1), (u2, v2), …, (uk, vk) with the smallest sums.

nums1 和nums2都是有序的，找最小sum的pair

输入: 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]
 

思路

可以通过大顶堆，并为每一个pair构造数组，1,2存num，3存sum

并且将堆大小保持在k。此时堆内所有数组节点都是需要的结果。

public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<List<Integer>> pq = new PriorityQueue<List<Integer>>((o1, o2)->o2.get(2)-o1.get(2));
        List<List<Integer>> ans = new ArrayList<>();
        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length; j++) {
                int sum = nums1[i] +nums2[j];
                if(pq.size()<k||sum <pq.peek().get(2)){
                    pq.offer(new ArrayList<>(Arrays.asList(nums1[i],nums2[j],sum)));
                    if(pq.size()>k){
                        pq.poll();
                    }
                }
                else{
                    break;
                }

            }
        }
        while (pq.size()>k){
            pq.poll();
        }
        for (List<Integer> integers : pq) {
            List<Integer> item = new ArrayList<>();
            item.add(integers.get(0));
            item.add(integers.get(1));
            ans.add(item);
        }
        return  ans;
    }

Tag

heap