Medium | LeetCode 378. 有序矩阵中第 K 小的元素 | 优先队列 | 二分法

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

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

示例 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

解题思路

方法一： 优先队列

将二维数组看作是N个一维数组, 采用归并的方法, 如题 Hard | LeetCode 23. 合并K个升序链表 | 分治 | 优先队列 将每个数字及其在矩阵当中的坐标存入优先队列当中。通过优先队列出队K次, 即可得到第K个小值。

public int kthSmallest(int[][] matrix, int k) {
   Queue<int[]> queue = new PriorityQueue<>(Comparator.comparingInt(o -> o[0]));
   int row = matrix.length;
   int colume = matrix[0].length;
   // 将每行的第一个元素加入优先队列
   for (int i = 0; i < row; i++) {
       queue.offer(new int[]{matrix[i][0], i, 0});
   }
   int[] heapTop = null;
   // 优先队列出队K次, 得到第K小数
   for (int i = 0; i < k; i++) {
       heapTop = queue.poll();
       if (heapTop[2] < colume - 1) {
           // 将出队元素的后面一个元素加入到优先队列当中
           queue.offer(new int[]{matrix[heapTop[1]][heapTop[2] + 1], heapTop[1], heapTop[2] + 1});
       }
   }
   return heapTop[0];
}

方法二： 二分法

矩阵最小值是matrix[0][0], 设为l,矩阵最大值为matrix[matrix.length-1][matrix[0].length-1] 设为r, 分别作为矩阵第K小值的上界和下界。

然后从上界和下界取一个值, 为target。然后从矩阵中找小于等于target值的数量(这里可以把二维矩阵看做N个一维数组, 从这N个一维数组找小于等于target值的数量, 然后求和)。

如果小于等于target值的数字的数量小于K, 则说明这个target值一定过小了。

如果小于等于target值的数字的数量大于K, 这个target值可能过大了, 当然也有可能不过大, 因为target值在矩阵当中可能不止一个。

如果小于等于target值的数字的数量等于K, 这个target值可能刚好合适, 也有可能过大了。

这样二分的规则就出来了。

public int kthSmallest(int[][] matrix, int k) {
    int l = matrix[0][0], r = matrix[matrix.length-1][matrix[0].length-1];
    while (l < r) {
        int mid = (l + r) >> 1;
        if (check(matrix, k, mid)) {
            l = mid + 1;
        } else {
            r = mid;
        }
    }
    return l;
}

public boolean check(int[][] matrix, int k, int target) {
    int sum = 0;
    // 注意这里要求的是小于等于target值的数量, 也就是找大于target值的下标
    // r的值必须是, matrix[0].length 而不能是其值减1
    int l = 0, r = matrix[0].length; 
    for (int i = 0; i < matrix.length; i++) {
        // 每次找当前行时, 二分左边界还是0, 但是右边界一定不大于上次定位的r
        // 因为同一列, 下一行的值一定大于上一行
        l = 0;        
        while (l < r) {
            int mid = (l + r) >> 1;
            if (matrix[i][mid] <= target) {
                l = mid + 1;
            } else {
                r = mid;
            }
        }
        sum += l;
    }
    // 如果矩阵中小于等于target的数量 小于K, 则说明target值过小了
    // 如果矩阵中小于等于target的数量 大于K, 则说明target值可能存在过大了,
    // 但也有可能刚好是这个target值, 因为target值可能有重复
    return sum < k;
}