[LeetCode] 373. Find K Pairs with Smallest Sums

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.

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]

Constraints:

• 1 <= nums1.length, nums2.length <= 105
• -109 <= nums1[i], nums2[i] <= 109
• nums1 and nums2 both are sorted in ascending order.
• 1 <= k <= 104

查找和最小的 K 对数字。

给定两个以升序排列的整形数组 nums1 和 nums2, 以及一个整数 k。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。

找到和最小的 k 对数字 (u1,v1), (u2,v2) ... (uk,vk)。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

因为 input 数组是有序的，所以直觉上你会往双指针的思路上靠，但是后来发觉是不行的。为什么不行呢？看第二个例子，一开始你肯定是两个数组各自的第一个元素拿出来组成第一个 pair，但是之后呢，你怎么决定到底要移动哪个数组的指针呢？并不是说因着你用过了 nums1[i] 你就需要移动 i，同理，并不是说你用过了 nums2[j] 你就要移动 j。同时注意 i 指针和 j 指针可以指向同一个下标。可以想象一个极端的例子，如果 nums1 的第一个元素非常非常小，然后 nums2 里面的元素都很大，很可能这 K 个 pair 里面都有 nums1 的这第一个元素。

正确的思路是用 priority queue 创建一个 pair 中两个数的和的最小堆，最后弹出 K 个 pair 即可。但是这个做法现在因为加了新的testcase，会超时。

时间O(nlogk)

空间O(n)

Java实现

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new ArrayList<>();
        // corner case
        if (nums1 == null || nums2 == null || nums1.length == 0 || nums2.length == 0) {
            return res;
        }

        // normal case
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> (a[0] + a[1]) - (b[0] + b[1]));
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            for (int j = 0; j < Math.min(k, nums2.length); j++) {
                queue.offer(new int[] { nums1[i], nums2[j] });
            }
        }
        while (k != 0 && !queue.isEmpty()) {
            int[] cur = queue.poll();
            List<Integer> list = new ArrayList<>();
            list.add(cur[0]);
            list.add(cur[1]);
            res.add(list);
            k--;
        }
        return res;
    }
}

 

优化的思路是改成一个最大堆。还是先把所有的 pair 放入 priority queue 中，当这个优先队列的 size 小于 K 的时候，直接加入即可；但是当这个优先队列的 size 大于等于 K 了，但是此时还有元素需要加进去的时候，比较一下堆顶元素的 sum 和要加进去的元素的 sum 谁更大，抛弃掉更大的那个元素。最后优先队列里剩下前 K 小的元素。将这些元素再加入结果集中输出即可。可惜这个做法现在因为加了新的testcase，也会超时。

时间O(nlogk)

空间O(n)

Java实现

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // corner case
        List<List<Integer>> res = new ArrayList<>();
        if (nums1 == null || nums2 == null || nums1.length == 0 || nums2.length == 0) {
            return res;
        }
        
        // normal case
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> (b[0] + b[1]) - (a[0] + a[1]));
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            for (int j = 0; j < Math.min(k, nums2.length); j++) {
                if (queue.size() < k) {
                    queue.offer(new int[] {nums1[i], nums2[j]});
                } else {
                    int sum = nums1[i] + nums2[j];
                    int[] peek = queue.peek();
                    int peekSum = peek[0] + peek[1];
                    if (sum < peekSum) {
                        queue.poll();
                        queue.offer(new int[] {nums1[i], nums2[j]});
                    }
                }
            }
        }
        while (k != 0 && !queue.isEmpty()) {
            List<Integer> list = new ArrayList<>();
            int[] cur = queue.poll();
            for (int num : cur) {
                list.add(num);
            }
            res.add(list);
            k--;
        }
        return res;
    }
}

 

[2022年8月更新] 这里我再提供一个不超时的方法，大体思路也是用最小堆实现归并排序，类似 23题 和 373题。注意题目给的条件，两个数组都是各自有序的，那么假如对于一个来自 nums1 的数字 u1 而言，他可以和 nums2 里面每个数字 v_x 都形成一个 pair，并且他们的和也是有序的。

[ [u1, v1], [u1, v2], [u1, v3], [u1, v4], ...[u1, v_n-1], [u1, v_n] ]

那么对于在 nums1 里的每个数字，他都可以和 nums2 里所有的数字组成这样一个有序的数组。所以这个题目就可以转化为对多个有序的数组进行归并排序并找出其中最小的 K 个元素。如下这个矩阵里，每一行的 pair sum 是从小到大排列的，每一列也是从小到大排列的。

[ [u1, v1], [u1, v2], [u1, v3], [u1, v4], ...[u1, v_n-1], [u1, v_n] ]

[ [u2, v1], [u2, v2], [u2, v3], [u2, v4], ...[u2, v_n-1], [u2, v_n] ]

[ [u3, v1], [u3, v2], [u3, v3], [u3, v4], ...[u3, v_n-1], [u3, v_n] ]

……………………………………………………………………

[ [u_{m - 1}, v1], [u_{m - 1}, v2], [u_{m - 1}, v3], [u_{m - 1}, v4], ...[u_{m - 1}, v_n-1], [u_{m - 1}, v_n] ]

[ [u_m, v1], [u_m, v2], [u_m, v3], [u_m, v4], ...[u_m, v_n-1], [u_m, v_n] ]

 

所以具体的做法是创建一个 PriorityQueue，里面存的是来自 nums1 和 nums2 的下标。初始化的时候，我们把所有的来自 nums1 的下标与 0 打包，放到最小堆里，模拟的是矩阵里第一列的内容。此时 PriorityQueue 里有 K 个元素，且堆顶元素是全局最小的。我们不断往 PriorityQueue 增添新的元素，直到结果集里有 K 个结果为止。

时间O(nlogn)

空间O(n)

Java实现

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new ArrayList<>();
        PriorityQueue<int[]> queue = new PriorityQueue<>(
                (a, b) -> (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]]));
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            queue.offer(new int[] { i, 0 });
        }

        while (!queue.isEmpty() && res.size() < k) {
            int[] cur = queue.poll();
            int a = cur[0];
            int b = cur[1];
            res.add(Arrays.asList(nums1[a], nums2[b]));
            if (b + 1 < nums2.length) {
                queue.offer(new int[] { a, b + 1 });
            }
        }
        return res;
    }
}

 

相关题目

373. Find K Pairs with Smallest Sums

378. Kth Smallest Element in a Sorted Matrix

719. Find K-th Smallest Pair Distance

LeetCode 题目总结