LeetCode 378 有序矩阵中第 K 小的元素：python3 题解

1. 题目理解输入一个 × 的矩阵matrix。一个整数k。关键特性矩阵的每一行从左到右是升序排列的。矩阵的每一列从上到下是升序排列的。目标找到矩阵中所有元素排序后的第 小元素注意如果有重复元素重复的也要算进去。例如[1, 1, 2]中第 2 小的是 1第 3 小的是 2。限制条件内存复杂度必须优于 (2)。这意味着我们不能简单地把所有元素拿出来排序因为那样需要 (2) 的空间存储所有元素。 最大为 300。2. 解题思路分析这道题主要有两种经典的解法分别利用了堆优先队列和二分查找。2.1 方法一最小堆优先队列核心思想既然每一行都是有序的那么每一行的第一个元素一定是该行最小的。我们可以把这 行的“当前最小元素”都放入一个最小堆中。初始时将每一行的第一个元素放入堆中。每次从堆中弹出最小的元素这就是当前全局未访问元素中的最小值。弹出后将该元素所在行的下一个元素放入堆中如果该行还有元素的话。重复上述操作 次第 次弹出的元素就是答案。图解假设矩阵如下k81 5 910 11 1312 13 15堆初始化[1, 10, 12](每行第一个)弹出 1 (第 1 小)推入 5。堆[5, 10, 12]弹出 5 (第 2 小)推入 9。堆[9, 10, 12]弹出 9 (第 3 小)该行无后续。堆[10, 12]弹出 10 (第 4 小)推入 11。堆[11, 12]弹出 11 (第 5 小)推入 13。堆[12, 13]弹出 12 (第 6 小)推入 13。堆[13, 13]弹出 13 (第 7 小)推入 15。堆[13, 15]弹出 13 (第 8 小) -返回 13复杂度分析时间复杂度(log⁡)。我们需要执行 次弹出操作每次堆调整需要 log⁡堆的大小最大为 。空间复杂度()。堆中最多存储 个元素每行一个。满足题目优于 (2) 的要求。2.2 方法二二分查找值域核心思想我们不是对“位置”进行二分而是对“数值范围”进行二分。矩阵中最小值是matrix[0][0]最大值是matrix[n-1][n-1]。答案一定在这个范围内。我们猜测一个中间值mid。统计矩阵中有多少个元素小于等于mid。如果数量 说明mid太小了答案在右半部分 (left mid 1)。如果数量 ≥说明mid可能是答案或者答案在左半部分 (right mid - 1)我们需要记录mid并继续尝试更小的值。当left right时记录的最后一次满足条件的mid即为答案。关键难点如何在 () 时间内统计小于等于mid的元素个数利用矩阵行列有序的特性我们可以从左下角开始搜索设当前位置为(row, col)初始为(n-1, 0)。如果matrix[row][col] mid说明当前元素及其上方的所有元素同一列都小于等于mid。这一列贡献了row 1个符合条件的元素。我们向右移动 (col 1) 去检查更大的数。如果matrix[row][col] mid说明当前元素太大了需要找更小的数。我们向上移动 (row - 1)。这样只需要走 2 步即可完成统计。复杂度分析时间复杂度(log⁡(max−min))。二分查找的次数取决于数值范围每次统计需要 ()。空间复杂度(1)。只需要几个变量不需要额外空间。这是最优的空间解法。3. 代码实现下面提供两种方法的 Python 3 代码。代码中包含了详细的注释。3.1 方法一最小堆实现from typing import Listimport heapqclass Solution:def kthSmallest(self, matrix: List[List[int]], k: int) - int:方法一最小堆 (Priority Queue)时间复杂度O(k * log n)空间复杂度O(n)n len(matrix)# 最小堆存储元组 (数值行索引列索引)# 初始化将每一行的第一个元素放入堆中min_heap []for r in range(n):# 推入 (值行号列号)heapq.heappush(min_heap, (matrix[r][0], r, 0))# 执行 k 次弹出操作# 第 k 次弹出的元素即为第 k 小的元素for _ in range(k):# 弹出当前堆中最小的元素val, r, c heapq.heappop(min_heap)# 如果这是第 k 次弹出直接返回# 注意循环是从 0 到 k-1所以当 _ k-1 时是第 k 次if _ k - 1:return val# 将该元素所在行的下一个元素推入堆中# 确保不越界if c 1 n:heapq.heappush(min_heap, (matrix[r][c 1], r, c 1))return -1 # 理论上不会执行到这里3.2 方法二二分查找实现【⭐】from typing import Listclass Solution:def kthSmallest(self, matrix: List[List[int]], k: int) - int:方法二二分查找 (Binary Search on Value)时间复杂度O(n * log(max - min))空间复杂度O(1)n len(matrix)# 确定二分查找的上下界# 最小值是矩阵左上角最大值是矩阵右下角left, right matrix[0][0], matrix[n - 1][n - 1]# 辅助函数统计矩阵中小于等于 target 的元素个数def countLessEqual(target: int) - int:count 0# 从左下角开始搜索row, col n - 1, 0# 只要还在矩阵范围内while row 0 and col n:if matrix[row][col] target:# 如果当前元素 target# 由于列是有序的当前元素上方的所有元素也都 target# 这一列共有 row 1 个元素符合条件 (索引 0 到 row)count (row 1)# 向右移动尝试更大的元素col 1else:# 如果当前元素 target# 说明当前元素太大了需要找更小的# 向上移动row - 1return count# 开始二分查找while left right:# 防止溢出的中间值计算方式mid left (right - left) // 2# 统计小于等于 mid 的元素个数count countLessEqual(mid)if count k:# 如果比 mid 小于等于的个数少于 k说明第 k 小的元素比 mid 大# 搜索右半部分left mid 1else:# 如果比 mid 小于等于的个数 k说明第 k 小的元素 mid# mid 有可能是答案也可能答案在左边# 搜索左半部分并保留 mid 作为潜在答案right mid# 当 left right 时循环结束left 即为答案return left4. 两种方法对比与总结特性方法一最小堆方法二二分查找核心逻辑多路归并思想每次取最小在数值范围内猜测并验证时间复杂度(log⁡)(log⁡(Range))空间复杂度()(1)适用场景 较小时效率很高 很大或追求极致空间时更优实现难度简单利用标准库中等需理解统计逻辑关于题目进阶问题恒定内存 (1)方法二二分查找满足了这一要求因为它只使用了几个变量没有使用与 相关的额外数据结构。() 时间复杂度这是一个非常高级的算法问题通常涉及 Frederickson Johnson 算法在普通面试中不要求实现。它利用了更复杂的矩阵选择逻辑。对于面试而言掌握二分查找解法通常已经是非常优秀的表现了。4.1 为什么不能直接 flatten扁平化排序虽然对于 300290,000直接sorted(sum(matrix, []))[k-1]在 LeetCode 上也能通过因为现代计算机处理 9 万个整数非常快但题目明确要求“内存复杂度优于 (2)。扁平化需要创建一个长度为 2 的新列表空间复杂度是 (2)。如果 扩大到 104扁平化会导致内存溢出MLE而上述两种方法依然可行。因此为了符合题目约束和应对更大规模数据不应使用扁平化排序。4.2 总结这道题是考察堆和二分查找在二维有序结构中应用的经典题目。如果想到的思路是多路归并可以使用最小堆。如果想到的是答案具有单调性如果 是答案那么比 大的数肯定也满足“至少有 个数小于它”可以使用二分查找。推荐在面试中优先展示二分查找解法因为它展示了更强的空间优化能力(1) 空间且代码量并不比堆解法多太多。