题目描述
给定一个 n x n 矩阵,其中每行和每列元素均按升序排序,找到矩阵中第k小的元素。
请注意,它是排序后的第k小元素,而不是第k个元素。
示例:
matrix = [
[ 1, 5, 9],
[10, 11, 13],
[12, 13, 15]
],
k = 8,
返回 13。
解题思路
和《剑指offer》第4题类似,但是不一样。其实这个题的“二分法”的思路跟《剑指offer》的面试题3的题目二思路更像。
-
- 找出二维矩阵中最小的数left,最大的数right,那么第k小的数必定在left~right之间。
- mid=(left+right) / 2;在二维矩阵中寻找小于等于mid的元素个数count。
- 若这个count小于k,表明第k小的数在右半部分且不包含mid,即left=mid+1, right=right,又保证了第k小的数在left~right之间。
- 若这个count大于k,表明第k小的数在左半部分且可能包含mid,即left=left, right=mid,又保证了第k小的数在left~right之间。
- 因为每次循环中都保证了第k小的数在left~right之间,当left==right时,第k小的数即被找出,等于right。
注:这里的left mid right是数值,不是索引位置。
参考代码
typedef long long LL;
class Solution
{
public:
int kthSmallest(vector<vector<int>>& matrix, int k){
int rows = matrix.size();
int cols = matrix[0].size();
LL left = matrix[0][0];
LL right = matrix[rows-1][cols-1];
while(left < right){
// 每次循环都保证第K小的数在left~right之间,当left==right,第k小的数就是left
LL mid = (left + right) >> 1;
// 找二维矩阵中<=mid的元素总个数
LL cnt = findNotBiggerThanMid(matrix, rows, cols, mid);
if(cnt < k)
// 第k小的数在右半部分,且不包含mid
left = mid + 1;
else
// 第k小的数在左半部分,可能包含mid
right = mid;
}
return left;
}
// 根据有序矩阵的性质,优化查找过程
int findNotBiggerThanMid(vector<vector<int>>& matrix, int rows, int cols, int mid){
// 从左下角开始找
int i = rows - 1;
int j = 0;
int cnt = 0;
while(i >= 0 && j < cols){
if(matrix[i][j] <= mid){
// 第j列有i+1个元素<=mid
cnt += (i+1);
j++;
}
else{
// 第i行目前的数大于mid,需要继续在当前列往上找
i--;
}
}
return cnt;
}
};
最大堆
class Solution {
public:
int kthSmallest(vector<vector<int>>& matrix, int k) {
priority_queue<int> pq;
for (int i = 0; i < matrix.size(); ++i) {
for (int j = 0; j < matrix[0].size(); ++j) {
pq.push(matrix[i][j]);
if (pq.size() > k) pq.pop();
}
}
return pq.top();
}
};