文章目录
摘要
在算法题中，经常会遇到"从大量组合里找到前 K 个最小值"的问题。
LeetCode 373 就是一个典型的例子：两个有序数组，要求找出和最小的 K 对数对。
看似组合爆炸（n*m 个数对），但其实可以利用最小堆（优先队列）高效求解。
本文会用 Swift 给出题解，并结合实际场景来理解这道题。
描述
我们有两个升序排列的数组 nums1 和 nums2，再给定一个整数 k。
任务是：从所有可能的数对 (u, v)（u 来自 nums1，v 来自 nums2）里，找到和最小的前 k 个。
举几个例子：
-
示例 1
txt输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3 输出: [[1,2],[1,4],[1,6]]我们只要前三个最小的数对，而不是把所有组合列出来。
-
示例 2
txt输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2 输出: [[1,1],[1,1]]
难点在于：两个数组长度可以到 10⁵，如果直接枚举所有组合，时间复杂度会是 O(n*m)，完全不可行。必须用更高效的方法。
题解答案
核心思路是：
- 两个数组是 有序的，所以小的组合一定来自前面的元素。
- 我们可以用一个最小堆（优先队列）来维护候选的数对，每次取出最小的，扩展下一个候选。
- 类似于"多路归并"的思路：把
(nums1[i], nums2[0])当作起点，逐步往右扩展。
这样就能避免生成所有组合，而是按需取出前 k 个。
题解代码分析
下面是 Swift 的解法：
swift
import Foundation
struct Pair: Comparable {
let sum: Int
let i: Int
let j: Int
static func < (lhs: Pair, rhs: Pair) -> Bool {
return lhs.sum < rhs.sum
}
}
class Solution {
func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
var result = [[Int]]()
if nums1.isEmpty || nums2.isEmpty || k == 0 {
return result
}
// 最小堆（优先队列）
var heap = [Pair]()
// 初始时，把 nums1 的前 k 个元素和 nums2[0] 配对
for i in 0..<min(nums1.count, k) {
heap.append(Pair(sum: nums1[i] + nums2[0], i: i, j: 0))
}
heap.sort(by: <)
var count = 0
while !heap.isEmpty && count < k {
let smallest = heap.removeFirst()
result.append([nums1[smallest.i], nums2[smallest.j]])
count += 1
// 扩展下一个元素：同一行，右移一格
if smallest.j + 1 < nums2.count {
let newPair = Pair(sum: nums1[smallest.i] + nums2[smallest.j + 1],
i: smallest.i,
j: smallest.j + 1)
// 插入保持堆序
var idx = heap.binarySearchInsertIndex(of: newPair)
heap.insert(newPair, at: idx)
}
}
return result
}
}
extension Array where Element: Comparable {
// 二分查找插入位置（保持有序）
func binarySearchInsertIndex(of element: Element) -> Int {
var left = 0, right = self.count
while left < right {
let mid = (left + right) / 2
if self[mid] < element {
left = mid + 1
} else {
right = mid
}
}
return left
}
}代码解析
- Pair 结构体 ：存储一个数对
(nums1[i], nums2[j])，以及它们的和sum，用来比较。 - 初始化堆 ：只需要考虑
nums1的前 k 个元素和nums2[0]的组合，因为数组是升序的，后面的一定不会比前面小。 - 逐步扩展 ：每次从堆里拿出最小的组合
(i, j)，然后把(i, j+1)插入堆。 - 二分插入：这里用二分查找保持数组有序，代替真正的优先队列。
如果项目里可以用 优先队列库，代码会更简洁。但标准 Swift 没有自带，这里用二分插入模拟。
示例测试及结果
我们来跑几个测试：
swift
let solution = Solution()
print(solution.kSmallestPairs([1,7,11], [2,4,6], 3))
// 结果: [[1,2],[1,4],[1,6]]
print(solution.kSmallestPairs([1,1,2], [1,2,3], 2))
// 结果: [[1,1],[1,1]]
print(solution.kSmallestPairs([1,2], [3], 3))
// 结果: [[1,3],[2,3]]运行结果：
txt
[[1, 2], [1, 4], [1, 6]]
[[1, 1], [1, 1]]
[[1, 3], [2, 3]]完全符合题意。
时间复杂度
- 初始化堆：
O(k log k)（前 k 个元素排序）。 - 每次取出一个最小元素，需要
O(log k)插入新的元素。 - 最多执行
k次。 - 总复杂度 O(k log k)。
相比暴力的 O(n*m)，效率提升巨大。
空间复杂度
- 堆里最多存放
k个元素。 - 额外空间复杂度：O(k)。
总结
这道题其实就是一个典型的 "多路合并 + 最小堆" 应用。
思路上和合并 K 个有序链表、合并区间问题有点类似，都是"每次取最小，再扩展候选"。
在实际场景中，这种方法可以应用在：
- 推荐系统：比如推荐用户最可能喜欢的前 K 个商品，候选空间很大，但我们只关心前 K 个。
- 数据库查询优化：找两个表的"最小连接结果"。
- 搜索排序：需要在多个有序结果中快速找到最优前 K 个。