[378][中等][堆][二分] 有序矩阵中第K小的元素

题目描述

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

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

示例：

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

返回 13。

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

解题思路

归并排序 + 堆

由于每一行都是排好序的数组, 可以使用归并排序的思想, 类似于链表的操作, 从每行的第一个数字开始, 结合最小堆.

• 首先将所有行的第一个数字入堆, 每行的第一个数字肯定是这一行最小的, 因此全局最小的一定在其中(这里一定是左上角的值为最小值)

• 此时堆顶的数字为最小值, 将其弹出, 然后将其所在行的下一个数字入堆, 如果这个数字已经是这一行最后一个数字了, 则无需找到新入堆数字, 堆大小-1

• 这样重复$k-1$次, 此时堆顶的元素就是整个矩阵中第$k$小的元素

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        heap = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(heap)

        for i in range(k - 1):
            num, r, c = heapq.heappop(heap)
            if c < n - 1:
                heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
        return heap[0][0]

时间复杂度为$O(n^2\log{n})$, 空间复杂度为$O(n)$

其实这道题是23. 合并K个排序链表的简化版, 对于链表, 没有要求同列数字随行递增, 也没有要求每行的数字个数相同, 仍然可以使用归并排序+堆的套路解答. 但这同时也说明我们没有用到这些有用信息, 时间复杂度还有提升的空间.

二分查找

二分查找的思路参考: 有序矩阵中第K小的元素

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        left, right = matrix[0][0], matrix[n - 1][n - 1]  # 既然用二分法, 用`left`和`right`命名最合适
        while left < right:
            mid = (left + right) // 2
            if self.is_left(matrix, mid, k, n):
                right = mid
            else:
                """
                这里count加1是因为: 这里判断是在右半边, 而mid等于left时是被判别为左边的, 因此mid肯定不是最终的结果
                """
                left = mid + 1
        return left

    @staticmethod
    def is_left(mat, mid, k, n):
        i, j = n - 1, 0
        count = 0
        while i >= 0 and j < n:
            if mat[i][j] > mid:
                i -= 1
            else:
                count += i + 1
                j += 1
        return count >= k

相关题目

• [23][困难][堆] 合并K个排序链表