LeetCode 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 。

问题分析

此题用二分法。给定二维矩阵matrix的左上角的数字一定是最小的，而右下角的数字一定是最大的，所以我们设二分法的初始查找范围就是在这二者之间。我们算出中间数mid后，调用自定义的search_count函数，它返回矩阵中小于等于mid的数有多少个，如果这个结果小于k，那说明我们要找的第k小的数一定比mid大，所以我们将 l 置为mid加1；反之如果这个结果大于等于k，那么将 r 置为mid。待二分查找结束后，返回 l 或 r 即可。然后我们来说一下search_count函数如何实现：充分利用二维矩阵的每行和每列元素均按升序排序的性质，我们从矩阵的左下角元素开始，如果它小于等于目标数，那么我们将res加上 i + 1，也就是把包括该元素的往上的一列都算进来了，然后将 j++；如果大于目标数，那么我们将 i--。最终当 i 或 j 不在矩阵范围内时跳出循环，返回res。

代码实现

class Solution {
public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int m = matrix.size(), n = matrix[0].size();
        int l = matrix[0][0];
        int r = matrix[m-1][n-1];
        while(l < r){
            int mid = l + r >> 1;
            if(search_count(matrix, mid) < k)
                l = mid + 1;
            else
                r = mid;
        }
        return l;
    }
    
    int search_count(vector<vector<int>>& matrix, int& mid){
        int res = 0;
        int m = matrix.size(), n = matrix[0].size();
        int i = m - 1;
        int j = 0;
        while(i >= 0 && j < n){
            if(matrix[i][j] <= mid){
                res += i + 1;
                j++;
            }
            else
                i--;
        }
        return res;
    }
};