LeetCode 373. 查找和最小的 K 对数字
题目描述


题意分析
从两个非递减数组中各取一个元素组成数对,找出数对和最小的前
k对,返回每对中的两个数值。一次选择由两个数组下标共同确定,同一个元素可以与另一数组中的多个元素配对,并不是每个元素只能使用一次。数组可能含重复值。来自不同下标组合的数对即使数值相同,仍然是不同的选择,输出中要保留相应出现次数。和相等的候选不要求固定先后次序,只要选出的前
k对没有漏掉更小的和即可。
解法:小根堆多路归并
核心思路
[!blue]
直接生成所有数对需要处理
m × n种组合,即使只需要很少的结果也会浪费大量工作。利用数组有序的条件,固定nums1的下标i,让第二个下标依次取0、1、2…,这一行数对的和就按非递减顺序排列。问题因此变成从多条有序序列中取出最小的前k项。每一行只需要把尚未输出的第一项放进小根堆,状态保存数对和以及下标
(i, j)。行内其他项都不小于这一项,暂时无需入堆;所有行的首个候选中最小的堆顶,也就不大于任何尚未输出的数对,是全局下一项。弹出
(i, j)后,只把同一行的(i, j + 1)补入堆,继续维持“每条活动行一个最小未输出候选”的不变量。每个下标对只会通过前一列进入一次,所以不需要额外的访问集合;不同下标对应的相同值则应当照常输出。初始化也不必开启全部
m行,只需前min(m, k)行。如果某行下标不小于k,排在它前面的k个行首已经都不大于它的行首,而这一行后续数对更不会更小。因此前k个结果总能从保留的行里选出;遇到相同的和时,选择较早行中的候选仍然合法。得到
k对后停止。代码也检查堆是否为空,使输入为空、k为零或实际数对不足时能够自然结束,不会从空堆中取值。
解题步骤
- 任一数组为空或
k == 0时返回空列表。- 令
rows = min(nums1.length, k),把下标对(i, 0)的和与坐标加入小根堆,其中i遍历前rows行。- 弹出和最小的状态
(i, j),把对应的两个数组值加入答案。- 若这一行还有下一列,将
(i, j + 1)加入堆;其他行的候选保持不变。- 反复弹出并补入,直到已经输出
k对或堆变为空。
代码实现
class Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
List<List<Integer>> result = new ArrayList<>();
if (nums1.length == 0 || nums2.length == 0 || k == 0) {
return result;
}
PriorityQueue<long[]> heap = new PriorityQueue<>((a, b) -> Long.compare(a[0], b[0]));
// 后续行的首项前已有足够候选,只开启前 k 行
int rows = Math.min(nums1.length, k);
for (int i = 0; i < rows; i++) {
heap.offer(new long[] {
(long) nums1[i] + nums2[0],
i,
0
});
}
while (k-- > 0 && !heap.isEmpty()) {
long[] state = heap.poll();
int i = (int) state[1];
int j = (int) state[2];
result.add(Arrays.asList(nums1[i], nums2[j]));
// 只推进刚弹出的那一行,维持每行一个候选
if (j + 1 < nums2.length) {
heap.offer(new long[] {
(long) nums1[i] + nums2[j + 1],
i,
j + 1
});
}
}
return result;
}
}
import "container/heap"
type pair struct {
sum int64
i int
j int
}
type pairHeap []pair
func (h pairHeap) Len() int { return len(h) }
func (h pairHeap) Less(i, j int) bool { return h[i].sum < h[j].sum }
func (h pairHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *pairHeap) Push(x any) {
*h = append(*h, x.(pair))
}
func (h *pairHeap) Pop() any {
old := *h
item := old[len(old)-1]
*h = old[:len(old)-1]
return item
}
func kSmallestPairs(nums1, nums2 []int, k int) [][]int {
result := make([][]int, 0)
if len(nums1) == 0 || len(nums2) == 0 || k == 0 {
return result
}
// 后续行的首项前已有足够候选,只开启前 k 行
rows := len(nums1)
if k < rows {
rows = k
}
h := &pairHeap{}
heap.Init(h)
for i := 0; i < rows; i++ {
heap.Push(h, pair{
sum: int64(nums1[i]) + int64(nums2[0]),
i: i,
})
}
for k > 0 && h.Len() > 0 {
current := heap.Pop(h).(pair)
result = append(result, []int{
nums1[current.i],
nums2[current.j],
})
k--
// 只推进刚弹出的那一行,维持每行一个候选
if current.j+1 < len(nums2) {
nextJ := current.j + 1
heap.Push(h, pair{
sum: int64(nums1[current.i]) + int64(nums2[nextJ]),
i: current.i,
j: nextJ,
})
}
}
return result
}
复杂度分析
- 时间复杂度:
O((r + p) log(r + 1)),其中r = min(m, k)是初始化行数,p为实际输出数量。代码逐项入堆初始化,随后每次输出做一次弹出、至多一次补入,堆大小始终不超过r。- 空间复杂度:辅助空间
O(r),用于堆中的活动行候选;返回的p对结果另占O(p)。
关键点总结
[!green]
- 固定一个数组下标即可得到一条有序数对序列,堆只需比较各行尚未输出的第一项。
- 弹出哪一行就推进哪一行,保证覆盖完整且每个下标对只有一个生成来源。
- 第二个数组有序保证行内归并,第一个数组有序保证只初始化前
k行不会漏掉更小结果。- 排序依据是数对和,重复的值不代表重复的下标组合,不能按值去重。
易错点总结
[!yellow]
- 每次同时生成右侧和下一行候选:同一下标对可能从两个方向进入堆;本解法统一从行首开始,只向右推进。
- 推进其他行:弹出后必须补上同一行的下一列,否则会漏掉该行本应更早出现的候选。
- 把相同数值数对去重:不同下标组合仍是不同结果,去重会减少合法出现次数。
- 使用两数和相减作比较器:即使每个和能装入整数,两个和的差也可能溢出,代码使用宽整数与安全比较。
- 只看结果数量,不检查堆空:所有可用数对耗尽后需要停止,不能继续取堆顶。
相似题目
| 题目 | 难度 | 关联与区别 |
|---|---|---|
| 23. 合并 K 个升序链表 | 困难 | 同样做多路有序归并,本题每个固定的第一数组元素形成一条有序数对和序列。 |
| 378. 有序矩阵中第 K 小的元素 | 中等 | 两数组数对和形成行列有序矩阵,可对比求前k个对象与只求第k小数值的区别。 |
转载与许可
许可
本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议,转载请注明出处!