LeetCode 373. 查找和最小的 K 对数字
目录
题目描述
题意分析
给两个已经升序排好的整数数组
nums1和nums2，以及一个整数k。从第一个数组取一个数u、第二个数组取一个数v，就构成一个数对 $(u, v)$。要求返回和最小的k个数对。有几点约定要看清：数对是按下标取的，所以即便两数组里有重复值，不同下标组合也算不同的数对；返回的是数值对而不是下标对；若可组成的数对总数不足
k，就把所有数对都返回。约束信号：两个数组长度都可达 $10^5$，而
k不超过 $10^4$。数对总数最大是 $10^{10}$，绝不可能全部生成，但要输出的只有 $10^4$ 个——这个巨大落差说明算法的复杂度应该主要由k决定，而不是由数组长度决定。数组已排序则是一个必须利用的结构性条件。边界情况：任一数组为空或
k为 0 时返回空列表；k超过 $m \times n$ 时返回全部数对；元素可以是负数，和的范围是 $[-2 \times 10^9, 2 \times 10^9]$。
解法：小根堆多路归并
核心思路
问题关键： 固定
nums1[i]后，数对和
nums1[i] + nums2[0], nums1[i] + nums2[1], ...因为
nums2有序而构成一条非递减序列。问题等价于从多条有序序列中取最小的k个元素。为什么用小根堆： 堆中只保存每条序列“尚未取出的第一个数对”。堆顶就是所有未取数对中的最小者；弹出
(i,j)后，只需把同一行的下一个(i,j+1)入堆。不变量： 堆中每个活跃行恰有一个最小未取数对。已输出的是全局最小的若干数对，堆顶是下一候选。弹出并补入同一行后，不变量继续成立，因此按此顺序取出的前
k个数对正确。只需初始化前
min(m,k)行：若i >= k，那么前k行的首个数对之和都不大于第i行首项，它不可能严格排进前k；和相等时任选其中k个也满足题意。和与比较器使用 64 位，避免边界数据相加或相减溢出。
解题步骤
- 任一数组为空或
k = 0时返回空结果。- 将前
min(nums1.length, k)行的首项(i,0)放入小根堆。- 重复至多
k次：弹出堆顶，将对应数值对加入答案。- 若该行还有下一列，将
(i,j+1)入堆；堆空时提前结束。口述示例：
nums1 = [1,7,11]、nums2 = [2,4,6]、k = 3。堆先放入三行首项(1,2)、(7,2)、(11,2)；弹出(1,2)后补(1,4)，再依次弹出(1,4)、(1,6)，得到和最小的三对。边界： 不计算
m * n来限制循环，可避免乘法溢出；当k超过数对总数时，堆会自然变空。
代码实现
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
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]));
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
}
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
}
复杂度分析
设
r = min(m,k)，实际返回p个数对。
- 时间复杂度：$O((r+p)\log r)$；通常写作 $O(k\log\min(m,k))$。
- 空间复杂度：$O(r)$，不计返回结果；堆中每个活跃行最多一个状态。
关键点总结
- 将二维数对矩阵按行看成多条有序序列，再做堆式多路归并。
- 堆状态必须保存行列下标，弹出后才能只推进对应行。
- 初始化前
min(m,k)行即可，避免堆随数组长度无谓增大。- 比较数对和时使用 64 位或安全比较器，不要用两个和直接相减。
易错点总结
- 生成全部
m * n个数对再排序：数据规模下会超时并耗尽内存。- 弹出
(i,j)后把(i+1,j)入堆：会重复行、遗漏同一行的后续候选。- 所有
m行都入堆：结果正确，但当m很大、k很小时浪费空间。- 比较器写成
sumA - sumB：差值可能溢出并破坏堆序。- 忽略空数组或
k超过总数：前者访问越界，后者应在堆空时自然结束。
相似题目
| 题目 | 难度 | 考察点 |
|---|---|---|
| LCR 061. 查找和最小的 K 对数字 | 中等 | 与本题同题，可直接复用多路归并模板 |
| 264. 丑数 II | 中等 | 三条序列共享同一份历史结果，需要去重 |
| 378. 有序矩阵中第 K 小的元素 | 中等 | 只求第 $k$ 小，可改用二分值域加逐行计数 |
| 786. 第 K 个最小的质数分数 | 中等 | 比较键是分数，需避免浮点误差或改用二分 |
| 23. 合并 K 个升序链表 | 困难 | 序列以链表形式给出，且要求全部归并而非前 $k$ |
| 632. 最小区间 | 困难 | 归并过程中同时维护堆内最大值以刻画区间宽度 |