LeetCode378之有序矩阵中第 K 小的元素（相关话题：优先队列，二分）

目录

题目描述

解法一、优先队列

解法二、二分法

Java代码

Python代码

参考文章

---

题目描述

给你一个 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

解法一、优先队列

• 要找第k小的元素，一种最常规的做法就是使用优先队列。
• 找第k小的元素，就保留k个最小的元素，其中最大的那个就是答案，所以可以使用最大优先队列。
• 遍历矩阵中的元素，将元素添加到队列中，如果队列中元素数目MaxPQ.size() > k，就将堆点最大的元素弹出。
• 遍历结束后弹出堆顶元素，就是最小的k个元素中最大的，即第k小的元素。

这里可以利用矩阵的有序性做一点小的优化：

如果在遍历的过程中，队列中的元素数目已经为k了，且如果当前元素大于堆顶元素，这个元素放入队列中还会被弹出，所以就没必要放入。

并且遍历的内循环是从某一行的从左到右遍历，当前元素的右边元素比当前元素更大，也没必要放入队列，所以当MaxPQ.size() == k && num > MaxPQ.peek()，直接打断内循环，进行下一行的遍历。时间复杂度为O(n2log(k))，空间复杂度为O(k)

代码如下

public int kthSmallest(int[][] matrix, int k) {
		
        //倒序构造大顶堆
        PriorityQueue MaxPQ = new PriorityQueue<>(Collections.reverseOrder());
		for (int[] row : matrix) {
			for (int num : row) {
				if (MaxPQ.size() == k && num > MaxPQ.peek())
					break;
				MaxPQ.add(num);
				if (MaxPQ.size() > k)
					MaxPQ.remove();
			}
		}
		return MaxPQ.remove();
}
	

解法二、二分法

由题目给出的性质可知，这个矩阵内的元素是从左上到右下递增的（假设矩阵左上角为 matrix[0][0]）。

我们知道整个二维数组中 matrix[0][0] 为最小值，matrix[n - 1][n - 1]为最大值，现在我们将其分别记作 l 和 r。

可以发现一个性质：任取一个数 mid 满足 ，那么矩阵中不大于 mid 的数，肯定全部分布在矩阵的左上角。

例如下图，取 mid=8

我们可以看到，矩阵中大于 mid 的数就和不大于 mid 的数分别形成了两个板块，沿着一条锯齿线将这个矩形分开。其中左上角板块的大小即为矩阵中不大于 mid 的数的数量。

Java代码

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

        // 使用二分查找来定位第 k 小的元素
        while (left < right) {
            int mid = left + ((right - left) /2); // 计算中间值

            // 判断矩阵中小于等于 mid 的元素数量是否满足条件
            if (check(matrix, mid, k, n)) {
                right = mid; // 如果数量大于等于 k，则可能的值在左侧或就是 mid
            } else {
                left = mid + 1; // 如果数量小于 k，则可能的值在右侧
            }
        }
        return left; // left 和 right 相遇，找到了第 k 小的元素
    }

    public boolean check(int[][] matrix, int mid, int k, int n) {
        int i = n - 1; // 从左下角开始
        int j = 0;
        int num = 0; // 用于记录小于等于 mid 的元素数量

        // 遍历矩阵
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= mid) {
                // 当前元素小于等于 mid
                num += (i + 1); // 所有当前列中上方的元素都不大于 mid
                j++; // 向右移动
            } else {
                i--; // 向上移动
            }
        }
        return num >= k; // 返回是否有至少 k 个元素小于等于 mid
    }
}

Python代码

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        # 初始化二分查找的边界
        left, right = matrix[0][0], matrix[-1][-1]
        n = len(matrix)


        #这个函数用于统计矩阵中有多少个元素小于或等于给定的 num。
        #它从矩阵的左下角开始，根据当前元素与 num 的比较来决定是向上移动还是向右移动。
        #count 记录小于或等于 num 的元素个数。
        def check(num):
            count = 0
            i, j = 0, n - 1
            while i < n and j >= 0:
                if num >= matrix[i][j]:
                    # 如果num大于等于当前元素，增加计数
                    count += (j + 1)
                    i += 1  # 移动到下一行
                else:
                    j -= 1  # 移动到前一列
            return count

   
        # 算法不断将查找范围划分为两半，每次基于 check 函数的返回值更新 left 或 right。
        # 如果 mid 值（left 和 right 的中点）使得 check(mid) 的返回值大于或等于 k，
        # 则 mid 可能是我们要找的元素，或者答案在 mid 左边，
        # 所以把 right 更新为 mid。
        # 否则，意味着 mid 太小，我们需要在更大的数中查找，因此把 left 更新为 mid + 1。
        while left < right:
            mid = left + (right - left) // 2
            cnt = check(mid)
            if cnt >= k:
                right = mid  # 答案可能是mid或者在mid左边
            else:
                left = mid + 1  # 答案在mid右边

        # 循环结束时，left即为第k小的元素
        return left

解题总结

对比74. 搜索二维矩阵这道题不具备每行的第一个整数大于前一行的最后一个整数这个属性所以不能直接把二维矩阵转化为一维数据进行二分。而是直接对矩阵里的最大值和最小值进行二分。而且最后二分的值并不能保证一定落在矩阵中题目本身逻辑存在缺陷，由于特殊的测试数据的这个bug没有暴露。另外本题可以利用分治算法把时间复杂度降低到O(n)可以参考附件资源（内容超前仅仅做兴趣了解）。

参考文章

378.java 二分法（图解）/优先队列 两种方法详解

归并排序