[LeetCode] 373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in non-decreasing 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]

Constraints:
1 <= nums1.length, nums2.length <= 105
-109 <= nums1[i], nums2[i] <= 109
nums1 and nums2 both are sorted in non-decreasing order.
1 <= k <= 104
k <= nums1.length * nums2.length

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

正确的思路是用 priority queue 创建一个 pair 中两个数的和的最小堆，最后弹出 K 个 pair 即可。但是这个做法相当于是两层 for 循环的暴力解，会超时。

复杂度

时间O(nlogk)
空间O(n)

代码

Java实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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 和 378 题。注意题目给的条件，两个数组都是各自有序的，那么假如对于一个来自 nums1 的数字 u1 而言，他可以和 nums2 里面每个数字 vx 都形成一个 pair，并且他们的和也是有序的。

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

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

1
2
3
4
5
6
7
8
9
10
11
[ [u1, v1], [u1, v2], [u1, v3], [u1, v4], ...[u1, vn-1], [u1, vn] ]

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

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

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

[ [um - 1, v1], [um - 1, v2], [um - 1, v3], [um - 1, v4], ...[um - 1, vn-1], [um - 1, vn] ]

[ [um, v1], [um, v2], [um, v3], [um, v4], ...[um, vn-1], [um, vn] ]

我参考了这个帖子。具体的做法是创建一个最小堆，里面存的是[i, j, nums1[i] + nums[j]]。初始化的时候，我们把所有的来自 nums1 的下标与 0 打包，放到最小堆里，模拟的是矩阵里第一列的内容。

全局最小的元素应该是 [u1, v1]，这个没什么疑问。但是最小堆中下一个元素是什么呢？有可能是 [u1, v2]，也有可能是[u2, v1]，但是这会导致一个问题：当我们拿到 [u1, v2] 的时候，我们会把 [u2, v2] 入堆；当[u2, v1] 出堆的时候，[u2, v2]会被再次入堆，这里需要去重。

如何做去重的工作呢？可以用 hashset，但是帖子里提供了一个很巧妙的思路。我们先将所有的 [u1, v1]、[u2, v1]、[u3, v1] 等等都入堆，然后每次出堆的时候，只增加 v 的值，比如当我们从最小堆中弹出 [u1, v1] 的时候，入堆的元素是[u1, v2]，而不是 [u2, v1]。这样就可以做到去重。

复杂度

时间O(klogk)
空间O(n)

代码

Java实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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) -> a[0] - b[0]);
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            queue.offer(new int[] { nums1[i] + nums2[0], i, 0 });
        }
        while (res.size() < k) {
            int[] cur = queue.poll();
            int i = cur[1];
            int j = cur[2];
            res.add(Arrays.asList(nums1[i], nums2[j]));
            if (j + 1 < nums2.length) {
                queue.offer(new int[] { nums1[i] + nums2[j + 1], i, j + 1 });
            }
        }
        return res;
    }
}

相关题目

1
2
3
373. Find K Pairs with Smallest Sums
378. Kth Smallest Element in a Sorted Matrix
719. Find K-th Smallest Pair Distance
#leetcode #java #array #greedy #heap
[LeetCode] 373. Find K Pairs with Smallest Sums
https://shurui91.github.io/posts/1953746004.html
Author
Aaron Liu
Posted on
July 31, 2020
Licensed under