LeetCode 373中等堆
查找和最小的 K 对数字 图解题解
这道题到底在问什么
给定两个以非递减顺序排列的整数数组 nums1 和 nums2,以及一个整数 k。定义一对值 (u, v),其中 u 来自 nums1,v 来自 nums2。请找到和最小的 k 个数对。nums1 和 nums2 的长度都可能达到 10^5,k 可能达到 10^4,且 LeetCode 保证 k <= nums1.length * nums2.length。
- 输入
- nums1 = [1,7,11], nums2 = [2,4,6], k = 3
- 输出
- [[1,2],[1,4],[1,6]]
- 解释
- 按和从小到大看,前几个数对是 [1,2]、[1,4]、[1,6]、[7,2] ...
- 输入
- nums1 = [1,1,2], nums2 = [1,2,3], k = 2
- 输出
- [[1,1],[1,1]]
先想最直接的笨办法
暴力法浪费在“生成了很多根本不会进前 k 的大数对”。我们只需要始终知道当前最小的候选是谁。(动画第 3 步)
最优解:一步一步想明白
- 3暴力法浪费在“生成了很多根本不会进前 k 的大数对”。我们只需要始终知道当前最小的候选是谁。
- 4这就是多路归并:堆里始终放每条链当前最小的未使用候选,堆顶就是全局下一个最小数对。
- 5sum(i,j)=nums1[i]+nums2[j];nums2 升序 -> 同一行 sum 递增每个堆节点显示 sum,节点下方标注 i、j 和真实数对。固定 nums1[0]=1 时,和是 3、5、7;固定 nums1[1]=7 时,和是 9、11、13。每行内部递增,所以一开始只需要拿每行第一个候选参赛。
- 6heapq.heappush(heap, (nums1[0] + nums2[0], 0, 0))堆元素存 (sum, i, j)。sum 用来排序,i 和 j 用来回到 nums1[i]、nums2[j] 拿具体数对。
- 7heapq.heappush(heap, (nums1[1] + nums2[0], 1, 0))行 1 的第一个候选是 [7,2],和为 9。堆顶仍然是和为 3 的 [1,2]。
- 8for i in range(min(k, len(nums1))): push (nums1[i] + nums2[0], i, 0)k=3,nums1 也只有 3 行,所以放入三条链的第一个候选。若 nums1 很长,也最多放 k 行:第 k 行及以后的首候选和,大于等于前 k 行首候选里的最大和,不会比这 k 个首候选更早进入前 k 个答案。
- 9_, i, j = heapq.heappop(heap); ans.append([nums1[i], nums2[j]])堆顶一定是当前所有链里最小的候选,所以 [1,2] 可以确定进入答案。画面里的堆已经移除了 sum=3 的根节点。
- 10heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))弹出 (i,j) 后,只推进同一行的 (i,j+1)。不能乱推进下一行,因为下一行的第一个候选已经在堆里了。
- 11_, i, j = heapq.heappop(heap); ans.append([nums1[i], nums2[j]])[1,4] 是当前最小候选。画面已经移除了 sum=5,剩下的堆顶变成 [7,2],它的和是 9。
- 12heapq.heappush(heap, (nums1[0] + nums2[2], 0, 2))因为 nums2 升序,行 0 的候选按 3、5、7 递增。弹出 5 后,7 成为这一行新的最小未使用候选。
- 13_, i, j = heapq.heappop(heap); ans.append([nums1[i], nums2[j]])这已经收集到 k=3 个答案。画面已经移除了 sum=7,堆里还剩 [7,2]、[11,2],但题目只要前 3 个。
- 14while heap and len(ans) < k: # False because len(ans) == k刚弹出的 (0,2) 已经是行 0 最后一格;nums2 长度为 3,下标范围是 0..2,所以行 0 耗尽,不再入堆。堆里虽然还剩行 1、行 2 的候选,但 len(ans)==k,while 条件不成立,循环结束。
- 17看到多个有序来源要取整体前 k 小,先想“每路放一个候选进小根堆”。
⚠️ 容易写错的地方
✗ 错:把所有 m*n 个数对都入堆
✓ 对:只把每行当前最小候选入堆
m、n 都可到 10^5,全量生成会爆。
✗ 错:弹出 (i,j) 后推进 (i+1,j)
✓ 对:推进同一行的 (i,j+1)
每个 nums1[i] 对应一条按 j 递增的链。
✗ 错:初始把 nums1 所有行都入堆
✓ 对:只放 min(k, len(nums1)) 行
答案只有 k 个,更多行的首候选不可能比已经放入的前 k 行更早全部用完。
✗ 错:用 set 去掉重复值数对
✓ 对:允许重复位置产生相同值对
示例 2 中两个 [1,1] 都要返回。
完整代码(Python)
Python
import heapq
class Solution:
def kSmallestPairs(self, nums1, nums2, k):
if not nums1 or not nums2 or k == 0:
return []
heap = []
for i in range(min(k, len(nums1))):
heapq.heappush(heap, (nums1[i] + nums2[0], i, 0))
ans = []
while heap and len(ans) < k:
_, i, j = heapq.heappop(heap)
ans.append([nums1[i], nums2[j]])
if j + 1 < len(nums2):
heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
return ans复杂度
时间复杂度
O(k log min(k, m))
m=len(nums1)。初始堆最多 min(k,m) 个元素,最多弹出 k 次,每次堆操作取对数。
空间复杂度
O(min(k, m))
堆里最多保存每条活跃配对链的一个候选。
看不够?换成动画再走一遍
上面的推演每一步都对应一帧动画。点开交互动画版,能一步步看着 查找和最小的 K 对数字 的数据怎么变、指针怎么走,还能切 Python / Java / C++ 跟着练。
面试官可能追问
这道题为什么用「堆」,换最直接的暴力解会差在哪?+
堆抓住了本题的结构特征,把暴力解里重复的工作省掉;暴力解通常要多嵌套一层枚举,数据一大就超时。具体对比见上文「暴力解及其卡点」与「最优解逐步推演」两节。
时间复杂度为什么是 undefined?怎么推出来的?+
按上文复杂度小节的推导,时间复杂度为 undefined。
本平台为独立第三方培训机构,与华为技术有限公司无任何关联;课程的服务内容与权益以购买协议为准,学习效果因个人情况而异。「华为 OD」「华为可信」等仅为对岗位与考试方向的客观描述,相关商标归各自权利人所有。