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

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

1 题目描述（Leetcode题目链接）

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

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

返回 13。

说明: 你可以假设 k 的值永远是有效的, $1\le k\le n^2$ 。

2 题解

  暴力解法合并成一维数组，用堆排序。更多关于堆的内容。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        l = []
        for li in matrix:
            l += li
        return heapq.nsmallest(k, l)[-1]

另外可以直接以数组每一行的第一个最小元素为基准建堆，弹出k个就好了。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        heap_size = len(matrix)
        res = 0

        def min_heapify(i, heap_size):
            l = 2*i + 1
            r = 2*i + 2
            if l < heap_size and matrix[l][0] < matrix[i][0]:
                smallests = l
            else:
                smallests = i
            if r < heap_size and matrix[r][0] < matrix[smallests][0]:
                smallests = r
            if smallests != i:
                matrix[i], matrix[smallests] = matrix[smallests], matrix[i]
                min_heapify(smallests, heap_size)

        while k:
            k -= 1
            res = matrix[0].pop(0)
            if not matrix[0]:
                matrix[heap_size-1], matrix[0] = matrix[0], matrix[heap_size-1]
                heap_size -= 1
            min_heapify(0, heap_size)

        return res