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

本文已参与「新人创作礼」活动，一起开启掘金创作之路。

题目来源：leetcode.cn/problems/kt…

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

思路

优先队列

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

• 通过优先队列存下所有数组当前还未遍历的最小元素，每次弹出队首元素，并将队首元素所在数组的下一个元素放入队列（如果还有下一个），这样弹出 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];
    }

二分

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

初始位置是矩阵的左下角

• 如果当前位置 [i, j] 的值小于等于 t，表示当前列元素都小于等于 t，统计值加上 j + 1，j++，即当前位置右移

• 如果当前位置 [i, j] 的值大于 t，需要将当前位置上移，即 i--

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

于是可以使用二分求矩阵中第 k 小的元素，初始左边界为左上角元素，右边界为右下角元素。每次求出矩阵中小于等于当前中间值的个数 count，根据 count 与 k 的大小更新左右边界，最终边界值即为所求值

	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;
    }