利用二分法解决 leetcode 378. Kth Smallest Element in a Sorted Matrix

问题简述

给定一个 $n \times n$ 的矩阵，矩阵中每行和每列的元素都按升序排列。给定一个 $k (k \in [1, n^2])$， 求再整个矩阵中按从小到大排序为 $k$ 的元素。
例如：

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

解题思路

该矩阵的行和列都是按从小到大的顺序排列的，这样的问题一般都是可以用二分法解决的。再来看，我们的输入 $k$，输出是第 $k$ 大的元素，不难发现，$k$ 越大，输出的元素也就越大。好了，决定用二分法了~
首先，我们以整个矩阵中的最小值为下边界（左上角元素），最大值为上边界（右下角元素）。然后用二分的方法猜过去，猜输出的结果是多少。只要满足猜的条件，我们就会把整个元素记录下来，然后看看有没有比这个元素小，并且满足猜的条件的元素。这个过程的时间复杂度为 $O(\ log(maxVal-minVal)\ )$，其中的 $minVal$ 为矩阵左上角的元素，$maxVal$ 为矩阵右下角的元素。
好了，接下来就是看看怎么猜了~
这里猜的时候又用到了一次二分法，其目的是找到每行中不大于猜的值 $val$ 的数有多少，然后算出整个矩阵中不大于 $val$ 的数有多少。如果不大于 $val$ 的数大于 $k$ 则返回true，否则返回false。这个过程的时间复杂度为 $O(\ nlog(n)\ )$。这一步是可以进一步优化的~
这个过程的时间复杂度为 $O(\ nlog(n)log(maxVal-minVal)\ )$。
最后，上代码~

代码

class Solution {
private:
    bool guess(vector<vector<int>>& matrix, int k, int val){
        for (int i = 0; i < matrix.size(); i++){
            int left = 0;
            int right = matrix[i].size() - 1;
            int ans = 0;
            while (left <= right){
                int mid = left + (right - left) / 2;
                if (matrix[i][mid] <= val){
                    left = mid + 1;
                    ans = mid + 1;
                }
                else{
                    right = mid - 1;
                }
            }
            k -= ans;
            if (k < 1)
                return true;
        }
        return false;
    }

public:
    int kthSmallest(vector<vector<int>>& matrix, int k) {
        int n = matrix.size();
        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];
        int ans;
        while (left <= right){
            int mid = left + (right - left) / 2;
            if (guess(matrix, k, mid)){
                ans = mid;
                right = mid - 1;
            }
            else{
                left = mid + 1;
            }
        }
        return ans;
    }
};