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

有序矩阵中第 K 小的元素

题目链接

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

解题思路

1. 在有序矩阵中，我们很容易的找到最小值 \(left\) 和最大值 \(right\) 。我们在 \([left, right]\) 范围中找进行二分，找出一个数 \(mid\)，有序矩阵中小于等于其的数有 \(K\) 个。
2. 如果有序矩阵中小于 \(mid\) 的数大于等于 \(K\) 个，说明我们选取的数 \(mid\) 大了，那么要找的数一定在\([left, mid - 1]\) 范围内，重新在此范围内二分。（我们要找的数一定在 \([left, mid - 1]\) 内）
3. 如果有序矩阵中小于 \(mid\) 的数小于 \(K\) 个，说明我们选取的数 \(mid\) 小了，那么要找的数一定在 \([mid + 1, right]\) 范围内，重新在此范围内二分。（我们要找的数一定在 \([mid + 1, right]\) 内）
4. 我们二分到最后的一个区间一定是 \([left, right]\)，其中 \(left = right\)。（我们要找的数一定在 \([left, right]\) 内）。那么我们要找的数就是 \(left 或者 right\)。
5. 循环结束后，\(left = right + 1\) 。根据循环不变量，在矩阵 \(matrix\) 中比区间 \([right + 1, ...]\) 中的数小的个数是大于等于 \(k\) 的。同理，在矩阵 \(matrix\) 中比区间 \([..., left - 1]\) 中的数小的个数是小于 \(k\) 的。
6. 我们要找矩阵中第 \(K\) 小的元素是 \(right + 1 = left\)。（在 \([left, right]\) 中找第一个数，在有序矩阵中比其小的个数是大于等于 \(K\) 的）

解题代码

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int l = matrix[0][0], r = matrix[n-1][n-1];

        while (l <= r) {
            int mid = l + (r - l) / 2;
            if (check(matrix, mid) >= k) {
                // 要寻找的数比 mid 小
                r = mid - 1;
            } else {
                // 要寻找的数比 mid 大
                l = mid + 1;
            }
        }
        return l;
    }

    // 计算矩阵 matrix 中小于等于 mid 的元素的个数
    public int check(int[][] matrix, int target) {
        // 总的思路是遍历 matrix 中的每一列中小于等于 mid 的数有多少个。
        int row = matrix.length;
        int i = row - 1, j = 0, cnt = 0;
        // 当行越界或者列越界，循环终止
        while (i >= 0 && j < row) {
            if (matrix[i][j] <= target) {
                // 当前列中的数都是小于 mid 的，因为矩阵中的元素在每一列上是升序排列的，因此，matrix[x][j], x <= i 均是小于等于 target 的，数量为 i + 1
                cnt += i + 1;
                // 判断下一列
                j++;
            } else {
                // 向上一行走。之前一直有疑问，为什么移动到了下一列，就不需要从最后一行从下往上判断了？
                // 假设在 matrix[i][j] 时 matrix[i][j] <= target ，需要移动到下一列进行判断。
                // 按道理讲是要从最后一行从下往上遍历的，即从 matrix[n - 1][j + 1] 开始进行判断，但是实际上
                // 只需要从 matrix[i][j + 1] 处进行判断即可，因为 matrix[i][j] 右下的元素都是大于等于 matrix[i][j] 的。 
                i--;
            }
        }
        return cnt;
    }
}

复杂度分析

时间复杂度：共需要进行 \(O(log(max - min))\) 次二分查找（即需要执行 \(O(log(max - min))\) 次 check()函数），\(max\) 是矩阵中的最大值，\(min\) 是矩阵中的最小值。在 check() 函数中，时间复杂度是 \(O(n)\)，\(n\) 为矩阵的周长。因此，总的时间复杂度为 \(O(nlog(max - min))\) 。

空间复杂度：\(O(1)\) 只用到了有限个数的变量。