有序矩阵中第K小的元素

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。

方法一：暴力

我们直接将矩阵中的元素存到一维数组，然后进行排序，得到第 K 小的元素值。

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int rows = matrix.length, cols = matrix[0].length;
        int index = 0;
        int[] sorted = new int[rows * cols];
        for(int[] mRow : matrix){
            for(int x : mRow){
                sorted[index++] = x;
            }
        }
        Arrays.sort(sorted);
        return sorted[k - 1];
    }
}

时间复杂度O(n^2 log n)，对 n^2 个数排序。

空间复杂度O(n^2)。

方法二：归并排序

使用小根堆维护，将最小的 k - 1 个值丢弃，最后的最小值就是答案。

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        //优先级队列
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>(new Comparator<int[]>() {
            public int compare(int[] a, int[] b) {
                return a[0] - b[0];
            }
        });
        int n = matrix.length;
        //将第一列的元素加入优先级队列
        for (int i = 0; i < n; i++) {
            pq.offer(new int[]{matrix[i][0], i, 0});
        }
        //依次取出最小的 k - 1 个数
        for (int i = 0; i < k - 1; i++) {
            int[] now = pq.poll();
            if (now[2] != n - 1) {
                pq.offer(new int[]{matrix[now[1]][now[2] + 1], now[1], now[2] + 1});
            }
        }
        return pq.poll()[0];
    }
}

时间复杂度O(k log n) 归并 k 次，每次调整堆 O(log n)。

空间复杂度O(n)，堆大小始终为 n。

方法三：二分查找

由题目给出的性质可知，这个矩阵内的元素是从左上到右下2递增的，在整个二维数组中 matrix[0][0] 为最小值， matrix[n - 1][n - 1] 为最大值，将其记作 l 和 r。我们发现：任取一个数 mid ，满足 l <= mid <= r，那么矩阵中不大于 mid 的数肯定分布在矩阵的左上角。

例如下图，取mid = 8：

我们可以看到，矩阵中大于 mid 的数和不大于 mid 的数被分成了两块，其中左上角板块的大小即为矩阵中不大于mid 的数的数量。

我们只需沿着这条锯齿线走一遍即可计算出这两个板块的大小：

初始位置在 matrix[n - 1][0] 左下角；
设当前位置为 matrix[i][j] 。若 matrix[i][j] <= mid，则将当前所在列的不大于 mid 的数的数量(即 i + 1)累加到答案中，并向右移动，否则向上移动。
不断移动直到走出格子为止。

设答案为 x ，已知l <= x <= r，这样就确定了二分答案的上下界。

每次对于猜测的答案 mid ，计算矩阵中有多少数不大于 mid：

数量不多于 k，那么 mid <= x；
数量大于 k，那么 x < mid。

class Solution {
    public int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];
        
        while(left < right){
            int mid = left + (right - left) / 2;
            if(check(matrix, mid, k, n)){
                //不大于 mid 的数大于 k 个，x <= mid
                right = mid;
            }else{
                //不大于 mid 的数小于 k 个，x > mid
                left = mid + 1;
            }
        }
        return left;
    }

    public boolean check(int[][] matrix, int mid, int k, int n){
        int i = n - 1;
        int j = 0;
        int num = 0;
        while(i >= 0 && j < n){
            if(matrix[i][j] <= mid){
                //将第 j 列中小于 mid 的数的个数加到总数
                num += i + 1;
                //右移
                j++;
            }else{
                //上移
                i--;
            }
        }
        //不大于 mid 的数是否大于等于 k 
        return num >= k;
    }
}

时间复杂度：O(n log(r - l))，二分查找进行次数为 O(log(r - l))，每次操作时间复杂度O(n)。空间复杂度：O(1)。

关于为什么返回 left 及left = mid + 1 是否存在于数组中：

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

假设 left = 14，那么 check 函数返回 true ，此时选择 right = mid，也就是选择右边界往左缩，即使 mid = 14 为可能的解，二分查找并不会结束。我们要找的是满足不大于 mid 的数的个数等于 k 的最左值，可以想像成寻找有序数组中某个重复出现数字第一次出现的位置，这个最左值一定存在于数组中，所以直接返回 left。