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

题目来源：https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/

大致题意：
给一个 n*n 的有序矩阵（每一行、每一列都是升序排列），一个整数 k，找出矩阵中第 k 小的元素

思路

优先队列

把矩阵的每一行想成一个有序数组，该题就变成了在 n 个有序数组中找第 k 小的元素

解题思路可以概括为

1. 通过优先队列存下所有数组当前还未遍历的最小元素，具体实现时，优先队列的元素可以是一个大小为 3 的数组，内容依次为 [元素值，元素所在行，元素所在列]
2. 初始时优先队列存储的内容为每行的第一个元素，以及对应元素的位置
3. 弹出队首元素，并将队首元素所在数组的下一个元素放入队列（如果还有下一个）。这样弹出 k - 1 次后，队首元素即为第 k 小的元素。

之所以每次弹出元素后，将弹出元素所在行的下一个元素加入队列，是因为下一个最小的元素要么为已经加入队列中的元素，要么为已弹出元素的下一个

代码：

public int kthSmallest_MergeSort(int[][] matrix, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> a[0] - b[0]);
        int n = matrix.length;
        // 存下所有行的第一个元素
        // 优先队列的元素是一个大小为 3 的数组，内容为 [元素值，元素所在行，元素所在列]
        for (int i = 0; i < n; i++) {
            queue.offer(new int[]{matrix[i][0], i, 0});
        }
        int idx = 1;
        // 弹出 k - 1 次
        while (idx++ < k) {
            int[] cur = queue.poll();
            // 若当前元素对应行还有后续元素，放入队列
            if (cur[2] < n) {
                queue.offer(new int[]{matrix[cur[1]][cur[2] + 1], cur[1], cur[2] + 1});
            }
        }
        // 当前队首元素值即为所求值
        return queue.peek()[0];
    }

二分

传统的二分的思路是每次比较后可以忽略掉当前查询区间范围的一半，这样就达到了查询优化的效果。

对于本题，可以使用类似二分的思想求矩阵中第 k 小的元素。

1. 初始左边界为左上角元素，右边界为右下角元素
2. 每次求出矩阵中小于等于当前中间值的个数 count，根据 count 与 k 的大小更新左右边界，最终边界值即为所求值

可以通过如下方法求有序矩阵小于等于 t 的元素个数：

初始位置是矩阵的左下角

• 如果当前位置 [i, j] 的值小于等于 t，表示当前列在当前位置上侧的元素都小于等于 t，统计当前列小于等于 t 的元素个数后，需要在更大的范围找，则将 j + 1，即当前位置右移
• 如果当前位置 [i, j] 的值大于 t，表示当前行在当前位置右侧元素都大于 t，需要在缩小查找的范围，则将 i - 1，即当前位置上移

重复上述操作，直至当前位置不再矩阵中

具体看代码

	int n;
    public int kthSmallest_BinarySearch(int[][] matrix, int k) {
        n = matrix.length;
        // 左右边界
        long l = matrix[0][0];
        long r = matrix[n - 1][n - 1];
        while (l < r) {
            long mid = (l + r) >> 1;
            // 求出矩阵中小于等于 mid 的数的个数
            int count = getCount(matrix, mid);
            // count 小于 k 表示第 k 小的元素值大于 mid，更新左边界
            if (count < k) {
                l = mid + 1;
            } else {    // 否则，表示第 k 小的元素值小于等于 mid，更新右边界
                r = mid;
            }
        }
        return (int) l;
    }

    // 求出矩阵中小于等于 mid 的数的个数
    public int getCount(int[][] matrix, long target) {
        int count = 0;
        int row = n - 1;
        int col = 0;
        while (row >= 0 && col < n) {
            // 当前位置值小于等于给定目标，向右移动
            // 并统计当前列小于等于给定目标的数的个数
            if (matrix[row][col] <= target) {
                count += row + 1;
                col++;
            } else {    // 当前位置值大于给定目标，向上移动
                row--;
            }
        }
        return count;
    }