题目来源: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;
}