题目描述

给定一个二维数组 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]

示例 2：

输入：matrix = [[-5]], k = 1
输出：-5

解法

方法一：暴力法

由于矩阵 matrix 已按行和按列都按非递减顺序排列，我们可以将整个矩阵转化为一个一维数组，并将其排序，找出第 k 小的元素即可。时间复杂度为 $O(n^2logn)$，其中 $n$ 是矩阵的边长。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        nums = []
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                nums.append(matrix[i][j])
        nums.sort()
        return nums[k-1]

方法二：归并排序

方法一的时间复杂度较高，主要是由于将整个矩阵转化为一个一维数组后排序的时间花费过大。对于一个有序的数组，我们可以使用归并排序的方法进行优化。

我们可以构造一个小根堆，将矩阵中的元素依次加入堆中，并维护一个元素数量为 $k$ 的堆。最后堆中的元素即为矩阵中的前 $k$ 小元素。

from typing import List
import heapq

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        m, n = len(matrix), len(matrix[0])
        pq = [(matrix[i][0], i, 0) for i in range(m)]
        heapq.heapify(pq)
        for i in range(k - 1):
            num, x, y = heapq.heappop(pq)
            if y != n - 1:
                heapq.heappush(pq, (matrix[x][y + 1], x, y + 1))
        return heapq.heappop(pq)[0]

方法三：二分查找

我们可以将矩阵中的每一行看成有序数组，进行二分查找。具体地，设矩阵中最小的数为 left，最大的数为 right，则中间数为 (left+right)//2。遍历整个矩阵，统计小于等于中间数的元素个数，若个数不足 k 个，则说明第 k 小的数在右半部分数组中，否则在左半部分数组中。

from typing import List

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        m, n = len(matrix), len(matrix[0])
        left, right = matrix[0][0], matrix[-1][-1]
        while left < right:
            mid = (left + right) // 2
            cnt, j = 0, n - 1
            for i in range(m):
                while j >= 0 and matrix[i][j] > mid:
                    j -= 1
                cnt += j + 1
            if cnt < k:
                left = mid + 1
            else:
                right = mid
        return left

总结

在应对二维数组排序的问题时，我们可以从以下几个方面考虑：

• 方法一：暴力法，将二维数组转化为一维数组后进行排序。
• 方法二：使用归并排序的方法构造一个小根堆进行优化。
• 方法三：通过二分查找的方法，计算出中间数，统计小于等于中间数的元素个数，并根据这个个数的大小不断缩小搜索范围，最终找到第 $k$ 小的数。

以上三种方法分别有不同的适用场景与时间复杂度，建议根据具体情况进行选择。