摘要
在算法题中,经常会遇到“从大量组合里找到前 K 个最小值”的问题。
LeetCode 373 就是一个典型的例子:两个有序数组,要求找出和最小的 K 对数对。
看似组合爆炸(n*m 个数对),但其实可以利用最小堆(优先队列)高效求解。
本文会用 Swift 给出题解,并结合实际场景来理解这道题。
描述
我们有两个升序排列的数组 nums1
和 nums2
,再给定一个整数 k
。
任务是:从所有可能的数对 (u, v)
(u
来自 nums1
,v
来自 nums2
)里,找到和最小的前 k
个。
举几个例子:
示例 1
输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3 输出: [[1,2],[1,4],[1,6]]
我们只要前三个最小的数对,而不是把所有组合列出来。
示例 2
输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2 输出: [[1,1],[1,1]]
难点在于:两个数组长度可以到 10⁵,如果直接枚举所有组合,时间复杂度会是 O(n*m),完全不可行。必须用更高效的方法。
题解答案
核心思路是:
- 两个数组是 有序的,所以小的组合一定来自前面的元素。
- 我们可以用一个最小堆(优先队列)来维护候选的数对,每次取出最小的,扩展下一个候选。
- 类似于“多路归并”的思路:把
(nums1[i], nums2[0])
当作起点,逐步往右扩展。
这样就能避免生成所有组合,而是按需取出前 k
个。
题解代码分析
下面是 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 没有自带,这里用二分插入模拟。
示例测试及结果
我们来跑几个测试:
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]]
运行结果:
[[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 个。