手记

LeetCode 378. 有序矩阵中第K小的元素 | Python

378. 有序矩阵中第K小的元素


题目


给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。

请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。

示例:


matrix = [

[ 1, 5, 9],

[10, 11, 13],

[12, 13, 15]

],

k = 8,

  

返回 13。

提示:

  • 你可以假设 k 的值永远是有效的,1 ≤ k ≤ n2 。

解题思路


先审题,题目中给出的是二维数组,而问题需要求得的是二维数组中第 k 小的元素。在这里,最直接的方法就是将二维数组转换为一维数组,对于转换后的一维数组进行升序排序,那么此时的一维数组中的第 k 个元素就是所求的结果。代码大致如下:


# python

def  kthSmallest(self, matrix: List[List[int]], k: int) -> int:

res = []

for row in matrix:

for ele in row:

res.append(ele)

res.sort()

return res[k-1]

在这里,代码并没有使用矩阵的特性。而是转换为一维数组进行求解。这个解法的时间复杂度为 O(n^2logn),即是对 n x n 个数进行排序。

因为上面的解法是将矩阵转换为一维数组,并不是在矩阵的基础上解决问题。下面使用二分查找,来尝试以在矩阵的前提下,对问题进行分析解答。

思路:二分查找

考虑使用二分查找的原因,因为题目中说明,矩阵中,每行每列元素都是升序排序的。

也就是说在题目给定的矩阵当中,元素由左上到右下是递增的,以示例为基础扩充展开分析:

示例:


matrix = [

[ 1, 5, 9],

[10, 11, 13],

[12, 13, 15]

],

k = 8

将上面的示例稍微扩充如下(仅为更方便展示二分查找方法的效果):


matrix = [

[ 1, 5, 9, 10, 11],

[10, 11, 13, 14, 15],

[12, 13, 15, 16, 17],

[13, 14, 16, 17, 18],

[14, 18, 22, 26, 30]

],

k = 8

将上面二维数组转换为下图:

我们假设左上角的元素为 matrix[0][0],右下角的元素为 matrix[n-1][n-1],根据题意以及上面的示图可知。matrix[0][0] 是二维数组中的最小值,matrix[n-1][n-1] 是最大值。

我们现在使用二分查找的方法来分析,设 matrix[0][0]matrix[n-1][n-1] 分别为 leftright

现在我们尝试取 mid(left <= mid <= right),可以发现在矩阵当中小于等于 mid 值的元素会分布在矩阵的左上部分,而大于 mid 值的元素则分布在矩阵的右下部分。例如下图所示,此时取 mid 为 15:

在这里大于 mid 和小于等于 mid 的元素分为两部分,沿着红色的分割线将两者分开。此时,我们就可以看到,大于 mid 和小于等于 mid 两者元素的个数分别有多少。

上面红色分割线划分的依据,在这里进行分析:

因为每行每列的元素都是升序排列的,我们前面的分析,元素需要与 mid 进行比较。例如,我们现在要求分布在左边部分的元素,也就是元素值小于等于 mid 的部分。

此时我们可以考虑从左下角开始出发,往右上角去找。这样能够快速收缩范围。具体的流程如下:

  • 以左下角为起始点,往右上角开始扩散

  • 如果当前位置的元素值小于等于 mid 值时,说明从当前位置开始往上的所有元素都小于等于 mid(因为每列升序排列),记录当前元素个数 count,注意维护更新。

  • 此时满足元素值小于等于 mid 值时,向右边移动再次比较。否则,就向上移动,寻找稍小的元素进行比较,直至移动到边界。

在这里,当取 mid 值后进行划分时,按照上面的方法,能够确定小于或等于 k 的个数,那么就会出现以下两种情况:

  • 当 count 大于或等于 k 时,那么可以确定要找的答案在 [left, mid] 这边;

  • 否则,答案在 [mid+1, right] 这个区间。

那么根据这个思路,循环直至找到答案。

关于二分查找方法执行流程,可看如下简略图示(若不太理解,可手画图帮忙理解):

具体的代码实现如下。

代码实现



class  Solution:

def  kthSmallest(self, matrix: List[List[int]], k: int) -> int:

def  search(mid, n):

# 从左下角开始往右上角找

i, j = n-1, 0

# 统计小于或等于 mid 的元素个数

count = 0

while i >= 0  and j < n:

if matrix[i][j] <= mid:

# 当此时元素小于等于 mid

# 那么该元素往上的所有元素都会小于或等于 mid

count += (i + 1)

# 向右移动,继续比较

j += 1

else:

# 当此时元素大于 mid 时,说明这个元素不包含在内,往上移动

i -= 1

# count 与 k 值比对,判断所求元素落在哪边

return count >= k

n = len(matrix)

  

left, right = matrix[0][0], matrix[n-1][n-1]

  

while left < right:

mid = left + (right - left) // 2

if search(mid, n):

# 当 count 大于等于 k 时,表明答案落在 [left, mid]

# 否则落在 [mid+1, right]

right = mid

else:

left = mid + 1

return left

实现结果



1人推荐
随时随地看视频
慕课网APP