378 有序矩阵中第 K 小的元素(多路归并、二分查找)
1. 问题描述:
给你一个 n x n 矩阵 matrix ,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。
请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。
示例 1:
输入:matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出:13
解释:矩阵中的元素为 [1,5,9,10,11,12,13,13,15],第 8 小元素是 13
示例 2:
输入:matrix = [[-5]], k = 1
输出:-5
提示:
n == matrix.length
n == matrix[i].length
1 <= n <= 300
-10 ^ 9 <= matrix[i][j] <= 10 ^ 9
题目数据保证 matrix 中的所有行和列都按非递减顺序排列
1 <= k <= n ^ 2
来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix
2. 思路分析:
① 分析题目可以知道矩阵的每一行都是升序的,我们需要在矩阵排序之后的所有数字中找到第k小的元素,很明显可以使用多路归并的思想解决,这一题与力扣373题是类似的。我们先将矩阵中每一行的第一个数字加入到小根堆中,因为需要加入每一行的下一个数字,所以还需要维护当前序列的数字对应的横坐标与纵坐标,因为使用的是python语言所以使用元组将当前序列的数字以及对应的横坐标与纵坐标作为元组的成员加入到小根堆中。然后在循环中进行操作,当k > 1并且堆不为空的时候执行循环,此时的堆顶元素肯定是堆中的最小值,当弹出堆顶元素之后需要将当前序列的下一个数字加入到堆中,当我们在堆中弹出k - 1个元素的时候此时的堆顶元素肯定是第k小的元素这个时候我们弹出堆顶元素返回矩阵中对应位置的答案即可。
② 除了①中多路归并的思路,我们可以利用矩阵中每一行与每一列都是升序的特点,使用二分查找的方法,l = matrix[0][0], r = matrix[-1][-1],判断矩阵中当前的大于等于mid的数目是否大于等于k,如果满足说明答案肯定是在mid或者mid左边,否则在mid右边,在计算矩阵中大于等于mid数目的时候可以从每一行末尾开始枚举,因为上一行中小于等于mid的位置肯定是比下一行的位置要靠后的,所以从后往前枚举计算每一行小于等于mid的数目即可。
3. 代码如下:
多路归并:
from typing import List
import heapq
class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
queue = list()
r, c = len(matrix[0]), len(matrix[0])
for i in range(r):
# 元组中封装所需的信息, 包括每一行对应序列的数字以及横坐标与纵坐标
heapq.heappush(queue, (matrix[i][0], i, 0))
# 弹出k - 1个数字最后堆顶元素对应矩阵的值就是答案
while k > 1 and queue:
poll = heapq.heappop(queue)
# 当前序列还有数字可以加入到堆中
if poll[2] + 1 < c:
heapq.heappush(queue, (matrix[poll[1]][poll[2] + 1], poll[1], poll[2] + 1))
# k的数目减1
k -= 1
res = heapq.heappop(queue)
return matrix[res[1]][res[2]]
二分查找:
from typing import List
class Solution:
# 使用二分的思路求解
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
l, r = matrix[0][0], matrix[-1][-1]
while l < r:
mid = l + r >> 1
j = len(matrix[0]) - 1
count = 0
# 计算每一行小于等于mid的数目
for i in range(len(matrix)):
while j >= 0 and matrix[i][j] > mid: j -= 1
count += j + 1
# 答案是左边
if count >= k:
r = mid
# 答案在右边
else:
l = mid + 1
return l