378. Kth Smallest Element in a Sorted Matrix 有序矩阵中第K小的元素

Title

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

示例：

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

返回 13。

提示：

你可以假设 k 的值永远是有效的，1 ≤ k ≤ n2 。

直接排序

Solve

最直接的做法是将这个二维数组另存为一维数组，并对该一维数组进行排序，最后返回这各一维数组中的第k个元素即为答案。

Code

    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        rec = sorted(sum(matrix, []))
        return rec[k - 1]

复杂度分析

时间复杂度：O(n²logn)，对 n² 个数排序。

空间复杂度：O(n²)，一维数组需要存储这 n² 个数。

二分查找

Solve

由题目给出的性质可知，这个矩阵内的元素是从左上到右下递增的，由此可知整个二维数组中 matrix[0][0] 为最小值，matrix[n−1][n−1] 为最大值，现在我们将其分别记作 l 和 r。

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

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

定义一种走法：

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

我们只要沿着这条锯齿线走一遍即可计算出这两个板块的大小，也自然就统计出了这个矩阵中不大于 mid 的数的个数了。

我们发现这样的走法时间复杂度为 O(n)，即我们可以线性计算对于任意一个 mid，矩阵中有多少数不小于它。这满足了二分答案的性质。

不妨假设答案为 x，那么可以知道 l ≤ x ≤ r，这样就确定了二分答案的上下界。

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

• 如果数量不多于 k，那么说明最终答案 x 不小于 mid；
• 如果数量少于 k，那么说明最终答案 x 大于 mid。

这样我们就可以计算出最终的结果 x 了。

Code

	def kthSmallest_binarySearch(self, matrix: List[List[int]], k: int) -> int:
        def check(middle):
            i, j, num = length - 1, 0, 0
            while i >= 0 and j < length:
                if matrix[i][j] <= middle:
                    num += i + 1
                    j += 1
                else:
                    i -= 1
            return num >= k

        length, left, right = len(matrix), matrix[0][0], matrix[-1][-1]
        while left < right:
            mid = (left + right) // 2
            if check(mid):
                right = mid
            else:
                left = mid + 1
        return left

复杂度分析

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

空间复杂度：O(1)。