给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第 k 小的元素。
请注意,它是排序后的第 k 小元素,而不是第 k 个不同的元素。
示例:
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,
返回 13。
提示:
你可以假设 k 的值永远是有效的,1 ≤ k ≤ n2 。
前置知识
二分查找
堆
公司
阿里
腾讯
字节
思路
显然用大顶堆可以解决,时间复杂度 Klogn,其中 n 为矩阵中总的数字个数。但是这种做法没有利用题目中 sorted matrix 的特点(横向和纵向均有序),因此不是一种好的做法.
一个巧妙的方法是二分法,我们分别从第一个和最后一个向中间进行扫描,并且计算出中间的数值与数组中的进行比较,可以通过计算中间值在这个数组中排多少位,然后得到比中间值小的或者大的数字有多少个,然后与 k 进行比较,如果比 k 小则说明中间值太小了,则向后移动,否则向前移动。
/*
* @lc app=leetcode id=378 lang=javascript
*
* [378] Kth Smallest Element in a Sorted Matrix
*/
function notGreaterCount(matrix, target) {
// 等价于在matrix 中搜索mid,搜索的过程中利用有序的性质记录比mid小的元素个数
// 我们选择左下角,作为开始元素
let curRow = 0;
// 多少列
const COL_COUNT = matrix[0].length;
// 最后一列的索引
const LAST_COL = COL_COUNT - 1;
let res = 0;
while (curRow < matrix.length) {
// 比较最后一列的数据和target的大小
if (matrix[curRow][LAST_COL] < target) {
res += COL_COUNT;
} else {
let i = COL_COUNT - 1;
while (i < COL_COUNT && matrix[curRow][i] > target) {
i--;
}
// 注意这里要加1
res += i + 1;
}
curRow++;
}
return res;
}
/**
* @param {number[][]} matrix
* @param {number} k
* @return {number}
*/
var kthSmallest = function (matrix, k) {
if (matrix.length < 1) return null;
let start = matrix[0][0];
let end = matrix[matrix.length - 1][matrix[0].length - 1];
while (start < end) {
const mid = start + ((end - start) >> 1);
const count = notGreaterCount(matrix, mid);
if (count < k) start = mid + 1;
else end = mid;
}
// 返回start,mid, end 都一样
return start;
};
Python3:
class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
n = len(matrix)
def check(mid):
row, col = n - 1, 0
num = 0
while row >= 0 and col < n:
# 增加 col
if matrix[row][col] <= mid:
num += row + 1
col += 1
# 减少 row
else:
row -= 1
return num >= k
left, right = matrix[0][0], matrix[-1][-1]
while left <= right:
mid = (left + right) // 2
if check(mid):
right = mid - 1
else:
left = mid + 1
return left
CPP Code:
class Solution {
public:
bool check(vector<vector<int>>& matrix, int mid, int k, int n) {
int row = n - 1;
int col = 0;
int num = 0;
while (row >= 0 && col < n) {
if (matrix[row][col] <= mid) {
num += i + 1;
col++;
} else {
row--;
}
}
return num >= k;
}
int kthSmallest(vector<vector<int>>& matrix, int k) {
int n = matrix.size();
int left = matrix[0][0];
int right = matrix[n - 1][n - 1];
while (left <= right) {
int mid = left + ((right - left) >> 1);
if (check(matrix, mid, k, n)) {
right = mid - 1;
} else {
left = mid + 1;
}
}
return left;
}
};