力扣 373. 查找和最小的K对数字

原创已于 2022-04-08 09:58:24 修改 · 268 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#leetcode #算法

于 2022-01-14 10:52:09 首次发布

优先队列专栏收录该内容

9 篇文章

订阅专栏

本文介绍了一种使用优先队列解决寻找两个升序数组中和最小的k个数对的方法。通过初始化特定数对进入队列，并不断更新队列来确保找到正确的答案。

题目来源：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/

大致题意：
给两个升序数组，将两个数组的元素两两成对，返回和最小的 k 个数对

思路

设数组为 a, b，m 为数组 a 的长度，n 为数组 b 的长度
首先可以知道 (a[0], b[0]) 为最小的数对，而若 m 大于 k 时，第 k 小的数对一定小于等与 a[k] + b[0]
题目规定 k 最大值为 1000，于是可以考虑使用优先队列将候选的数对放入，每次取出队首最小的数对即可

优先队列

如何保证最小的数对一定在优先队列中？
简单的办法就是把 a[0] ~ a[min(m, k)] 与 b[0] ~ b[min(n, k)] 组成的所有数对放入队列，然后依次取出 k 个
这样做消耗比较大，可以进行优化

已知

若最小的数对为 (a[i], b[j])
那么次小的数对只可能为 (a[i], b[j + 1]) 或 a[i + 1] 与 b[0] ~ b[min(n, k)] 组成的某个数对

于是可以

初始时将 a[0] ~ a[min(m, k)] 与 b[0] 组成的数对放入队列
每次出队后，出队数对为 ( a[x], b[y] )，只要 y + 1 小于 n，那么就可以将 ( a[x], b[y + 1] ) 放入队列。这样保证了队首元素一定为次小的数对
出队 k 次，或者队列为空时停止

public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 优先队列
        PriorityQueue<int[]> queue = new PriorityQueue<>(k, (a, b) -> {
            return nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]];
        });
        List<List<Integer>> ans = new ArrayList<>();
        int m = nums1.length;
        int n = nums2.length;
        // 初始时放入a[0] ~ a[min(m, k)] 与 b[0] 组成的数对
        for (int i = 0; i < Math.min(m, k); i++) {
            queue.offer(new int[]{i, 0});
        }
        // 当出队 k 次或者队列为空时中止循环
        while (k-- > 0 && !queue.isEmpty()) {
            // 取队首元素
            int[] idx = queue.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[idx[0]]);
            list.add(nums2[idx[1]]);
            ans.add(list);
            // 若 idx[1] + 1 小与 nums2 长度，将之后位置的元素与 idx[0] 组成数对
            if (idx[1] + 1 < n) {
                queue.offer(new int[]{idx[0], idx[1] + 1});
            }
        }

        return ans;
    }