LeetCode LCR 061. 查找和最小的 K 对数字
题目描述
题意分析
从两个非空升序数组中各取一个元素，返回和最小的
k个数对；数对总数不足k时全部返回。不同下标可以产生相同的数值对，这些数对仍分别计数，不能按值去重。本题两个数组长度各至多
10^4，k至多为 1000。与其枚举全部m*n个组合，更适合利用数组有序性，逐个生成当前最小的候选。
解法：小根堆多路归并数对
核心思路
[!blue]
固定
nums1的下标i，依次与nums2[0]、nums2[1]……配对，数对的和单调不降。把每个i看作一行，就得到多条天然有序的数对序列，可以用小根堆做多路归并。堆中保存下标对
(i, j)，每条参与归并且尚未取完的行只放一个代表，代表这行剩余数对中最小的那个。任何未取数对都不小于自己的行代表，因此堆顶就是所有参与行中当前最小的候选。弹出
(i, j)后，输出对应的两个数组值；若这一行还有下一列，就加入(i, j+1)。同一行只通过前一项生成后一项，每个下标对只有一个来源，不需要额外的访问集合，也不能再从列方向重复加入候选。初始只放前
r = min(m, k)行的行首。对下标不小于k的任意一行，它的所有数对都不小于前k行的行首，已有至少k个保留候选不比它差。因此忽略后续行仍能得到一组正确答案；相等时可以用保留的同和数对替代，不应断言被忽略的数对绝不可能出现在某个合法答案中。每次取出最小值并补上同行后继，候选状态会继续成立。取到
k项就停止；若堆先空，说明可用数对已经取完。Java 使用Long.compare比较两个和，Go 直接比较大小，都避免了用两和相减构造比较结果。
解题步骤
- 建立按
nums1[i] + nums2[j]排序的小根堆，将前min(m,k)行的(i,0)加入。- 只要堆不空且结果不足
k项，就弹出当前最小的下标对。- 将对应的两个值加入答案，再检查该行是否还有下一列；有则加入同行后继。
- 满足数量或堆耗尽时返回，保留来自不同下标的重复数值对。
代码实现
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> ans = new ArrayList<>();
PriorityQueue<int[]> pq =
new PriorityQueue<>(
(a, b) ->
Long.compare(
(long) nums1[a[0]] + nums2[a[1]],
(long) nums1[b[0]] + nums2[b[1]]));
for (int i = 0; i < nums1.length && i < k; i++) {
pq.offer(new int[] {
i,
0
});
}
while (!pq.isEmpty() && ans.size() < k) {
int[] cur = pq.poll();
ans.add(List.of(nums1[cur[0]], nums2[cur[1]]));
if (cur[1] + 1 < nums2.length) {
pq.offer(new int[] {
cur[0],
cur[1] + 1
});
}
}
return ans;
}
}
import (
"container/heap"
)
type pairHeap struct {
nums1 []int
nums2 []int
idx [][2]int
}
func (h pairHeap) Len() int { return len(h.idx) }
func (h pairHeap) Less(i, j int) bool {
a, b := h.idx[i], h.idx[j]
return h.nums1[a[0]]+h.nums2[a[1]] < h.nums1[b[0]]+h.nums2[b[1]]
}
func (h pairHeap) Swap(i, j int) { h.idx[i], h.idx[j] = h.idx[j], h.idx[i] }
func (h *pairHeap) Push(x any) { h.idx = append(h.idx, x.([2]int)) }
func (h *pairHeap) Pop() any {
n := len(h.idx)
v := h.idx[n-1]
h.idx = h.idx[:n-1]
return v
}
func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
h := &pairHeap{nums1: nums1, nums2: nums2}
for i := 0; i < len(nums1) && i < k; i++ {
h.idx = append(h.idx, [2]int{
i,
0,
})
}
heap.Init(h)
ans := make([][]int, 0, k)
for h.Len() > 0 && len(ans) < k {
cur := heap.Pop(h).([2]int)
ans = append(ans, []int{
nums1[cur[0]],
nums2[cur[1]],
})
if cur[1]+1 < len(nums2) {
heap.Push(h, [2]int{
cur[0],
cur[1] + 1,
})
}
}
return ans
}
复杂度分析
设数组长度分别为
m、n，参与归并的行数为r = min(m,k)，实际返回数对数为p = min(k,m*n)。
- 时间复杂度：$O((r+p)\log(r+1))$，初始化至多
r个候选，随后至多进行p轮弹出和补入，堆规模不超过r。Go 的批量建堆本身为 $O(r)$，也包含在这一上界内。- 空间复杂度：堆为 $O(r)$。Java 结果按实际输出增长，占 $O(p)$；Go 预分配结果容量为
k，占 $O(k)$，实际返回其中的p个数对。
关键点总结
[!green]
- 固定第一数组下标后，另一维产生有序序列，可以直接多路归并。
- 堆保存每行的最小未取项，弹出后只补同行后继。
- 前
k行已经提供k个不劣于后续行的候选，足以安全限制堆规模。- 数对按下标组合计数，数值相同不代表重复生成了同一个下标对。
易错点总结
[!yellow]
- 只存数值、不保留下标：无法确定弹出候选属于哪一行，也就无法生成同行后继。
- 同时补同行和下一行的邻居：在已经初始化各行行首的写法中，会让同一下标对从多个方向重复入堆。
- 按两个值去重：不同下标的同值组合也应分别计入答案。
- 不检查堆是否为空：数对不足
k时仍继续弹出，会访问不存在的候选。- 用两个 32 位数对和直接相减作比较器：单个和合法，不代表它们的差仍在 32 位范围内。
- 把行数截断理解为严格排除同和候选：相等时可能有多组合法答案，截断保证能保留一组最优选择即可。
相似题目
| 题目 | 难度 | 关联与区别 |
|---|---|---|
| 23. 合并 K 个升序链表 | 困难 | 同样做多路有序归并，本题每个固定的第一数组元素形成一条有序数对和序列。 |
| 378. 有序矩阵中第 K 小的元素 | 中等 | 两数组数对和形成行列有序矩阵，可对比求前k个对象与只求第k小数值的区别。 |
转载与许可
许可
本博客所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处！