378-有序矩阵中第K小的元素

LeetCode 378. Kth Smallest Element in a Sorted Matrix

LeetCode-378

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the $\rm k^{th}$ smallest element in the matrix.

Note that it is the k^th^ smallest element in the sorted order, not the k^th^ distinct element.

You must find a solution with a memory complexity better than $O(n^2)$.

Example 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 3
Output: 9
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 3th smallest number is 9

method：二分法

因为矩阵已经从左到右、从上到下排列，所以左上角是最小值，右下角是最大值，以此作为左边界l和右边界r，在得到mid后，计算矩阵中小于等于mid的数量cnt

• 如果cnt >= k，说明mid是可行解，但最优解可能在[l, mid]，r要往左走，所以r = mid
• 如果cnt < k，说明mid不是可行解，最优解在[mid+1, r]，所以l = mid + 1

计算cnt的方法：

• i遍历每一行，因为已经从左到右排序，所以j最后一个元素开始找，直到matrix[i][j] <= mid，此时左边的元素肯定也小于等于mid，所以直接cnt += j + 1，然后换到下一行
• 因为从上到下是递增的，所以j就不用退回最后一列，直接从当前位置往左找第一个小于等于mid就行

bool check(vector<vector<int>>& matrix, int mid, int k) {
    int cnt = 0;
    int i = 0, j = matrix[0].size() - 1;
    while (i < matrix.size() && j >= 0) {
        if (matrix[i][j] <= mid) {
            cnt += j + 1;   // 当前行小于等于mid的个数
            i++;    // 下一行
        }
        else j--;   // j往左找第一个小于等于mid
    }
    return cnt >= k;
}
int kthSmallest(vector<vector<int>>& matrix, int k) {
    int n = matrix.size();
    int l = matrix[0][0];
    int r = matrix[n - 1][n - 1];
    while (l < r) {
        int mid = (l + r) >> 1; // 这里要向下取整
        if (check(matrix, mid, k)) r = mid;
        else l = mid + 1;
    }
    return l;
}

计算小于等于mid的个数的时间复杂度是$O(n)$，二分的时间复杂度是$O(log(r-l))$，所以总时间复杂度为$O(nlog(r-l))$
空间复杂度：$O(1)$

注意：除以2是向零取整，右移一位是向下取整

int a = -5;
cout << (a / 2) << endl;    // -2，向零取整
cout << (a >> 1) << endl;   // -3，向下取整