题目
给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。
请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。
示例:
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,
返回 13。
提示:
你可以假设 k 的值永远是有效的,1 ≤ k ≤ n2 。
解题思路
无脑版:
把二维数组合并成一维,然后返回第k个元素。时间复杂度
o
(
n
2
)
o(n^2)
o(n2)
归并排序版:
合并n个有序数组,然后返回第k个元素。归并的部分和合并K个排序链表思路一致
二分查找版:
二分查找是看了题解才想到的。用low, high
分别代表当前矩阵中的最大值和最小值,计算出中值mid = (low + high) // 2
,然后找矩阵中<= mid
的数字有几个。若<= mid
的数字不足k个,则最终的答案应该比mid
大。若<= mid
的数字>=k
个,则最终答案应该<=mid
代码
无脑版:
class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
last_list = matrix[0]
for each_list in matrix[1:]:
index_1, index_2 = 0, 0
new_merge = []
while index_1 < len(last_list) and index_2 < len(each_list):
if last_list[index_1] < each_list[index_2]:
new_merge.append(last_list[index_1])
index_1 += 1
else:
new_merge.append(each_list[index_2])
index_2 += 1
while index_1 < len(last_list):
new_merge.append(last_list[index_1])
index_1 += 1
while index_2 < len(each_list):
new_merge.append(each_list[index_2])
index_2 += 1
last_list = new_merge
return last_list[k - 1]
归并版:
class Solution:
def merge_2list(self, list1: list, list2: list) -> list:
ret_list = []
index1, index2 = 0, 0
while index1 < len(list1) and index2 < len(list2):
if list1[index1] < list2[index2]:
ret_list.append(list1[index1])
index1 += 1
else:
ret_list.append(list2[index2])
index2 += 1
if index1 != len(list1):
ret_list += list1[index1:]
if index2 != len(list2):
ret_list += list2[index2:]
return ret_list
def merge_sort(self, matrix: List[List[int]]) -> list:
if not matrix:
return None
if len(matrix) == 1:
return matrix[0]
mid = len(matrix) >> 1
return self.merge_2list(self.merge_sort(matrix[:mid]), self.merge_sort(matrix[mid:]))
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
sorted_list = self.merge_sort(matrix)
return sorted_list[k - 1]
二分查找版:
class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
def not_enough(target: int, k: int) -> bool:
"""
True: # of numbers that are no larger than target < k
False: # of numbers that are no larger than target >= k
"""
i, j = 0, n - 1
res = 0
while 0 <= i < n and 0 <= j < n:
if matrix[i][j] <= target:
i += 1
res += j + 1
else:
j -= 1
return res < k
n = len(matrix)
left, right = matrix[0][0], matrix[n - 1][n - 1]
while left < right:
mid = (left + right) >> 1
if not_enough(mid, k):
left = mid + 1
else:
right = mid
return (left + right + 1) >> 1