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

持续创作，加速成长！这是我参与「掘金日新计划 · 6 月更文挑战」的第5天，点击查看活动详情

题目描述

给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。 请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。

你必须找到一个内存复杂度优于 O(n2) 的解决方案。  

示例 1：
输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13

示例 2：
输入：matrix = [[-5]], k = 1
输出：-5

提示：

• n == matrix.length
• n == matrix[i].length
• 1 <= n <= 300
• -109 <= matrix[i][j] <= 109
• 题目数据 保证 matrix 中的所有行和列都按 非递减顺序 排列
• 1 <= k <= n2

思路

可以采取二分查找的思路。由于每行和每列都是升序排列，那么左上角就是最小值，右下角就是最大值。这里使用二分法，关键是如何快速求出小于mid的数量。这里有个小技巧，可以从矩阵的右上角开始，如果右上角的数据比mid大，那么j就往前走，直到j走到0或者找到第一个比mid小的数据，由于每行都是升序排列，所以这一行前面的值就不需要比较了；接下来下一行，不用从最后一个数开始，直接从j位置开始即可，这样，我们在O(n)的时间复杂度内就获取了比mid小的数量。接下来就是按照二分的方法进行即可。

Java版本代码

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int ans = -1;
        int left = matrix[0][0], right = matrix[n-1][n-1];
        while (left <= right) {
            int mid = left + ((right - left)>>1);
            // 获取小于等于mid的元素数量
            int cnt = 0;
            for (int i = 0, j = n-1; i < n && j >= 0;) {
                if (matrix[i][j] <= mid) {
                    cnt += j + 1;
                    i++;
                } else {
                    j--;
                }
            }
            if (cnt >= k) {
                ans = mid;
                right = mid - 1;
            } else {
                left = mid + 1;
            }
        }
        return ans;
    }
}