题目描述
几天没打题,感觉脑子都是一团浆糊。。。。
- 升序:肯定得用这条件来优化复杂度
- 数对:用 int[2] 来表示
思路 && 代码
- 看了题解区彤哥的思路,代码其实就 8 行= =
- 堆初始化:按照[i, 0]的规律初始化堆
- 弹出:当前堆顶就是最小值,弹出[index1, index2]后,往后最有可能的数对是[index1, index2 + 1],放入队列中继续循环~
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
// 小顶堆,每次弹出最小的当前结果
PriorityQueue<int[]> heap = new PriorityQueue<>
((a, b) -> nums1[a[0]] + nums2[a[1]] - (nums1[b[0]] + nums2[b[1]])); // 策略模式
// 堆初始化:把 nums1 的所有索引入队,nums2 的索引初始时都是 0
// 优化:最多入队 k 个就可以了,因为提示中 k 的范围较小,这样可以提高效率
for (int i = 0; i < Math.min(k, nums1.length); i++) heap.offer(new int[] {i, 0});
List<List<Integer>> ans = new ArrayList<>();
// 最多弹出 k 次(获取 k 对数字)
while (k-- > 0 && !heap.isEmpty()) {
// 弹出一个
int[] pos = heap.poll();
ans.add(Arrays.asList(nums1[pos[0]], nums2[pos[1]])); // 每次都弹出一个当前的最小值
// 加入一个,将 index2 加 1 之后继续入队
if (++pos[1] < nums2.length) heap.offer(pos);
}
return ans;
}
}
- 无注释版
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
PriorityQueue<int[]> heap = 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++) heap.offer(new int[] {i, 0});
List<List<Integer>> ans = new ArrayList<>();
while (k-- > 0 && !heap.isEmpty()) {
int[] pos = heap.poll();
ans.add(Arrays.asList(nums1[pos[0]], nums2[pos[1]]));
if (++pos[1] < nums2.length) heap.offer(pos);
}
return ans;
}
}