leetcode378. 有序矩阵中第K小的元素（Python3）

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

给定一个 n x n 矩阵，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。

请注意，它是排序后的第 k 小元素，而不是第 k 个不同的元素。

示例：

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

返回 13。

提示：

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

---

方法一：直接排序

思路：

最简单粗暴的方法就是直接将matrix中所有的数据进行排序，然后返回第k小的即可。

排序的方法有很多，我这里直接引入heapq来进行堆排序。

排序的时间复杂度很高，不是最好的方法。

代码：

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        import heapq
        nums = []
        n = len(matrix)
        #将matrix中所有数据加入双向队列中
        for i in range(n):
            for j in range(n):
                heapq.heappush(nums,matrix[i][j])
        #弹出前k个最小数的数组
        ans = heapq.nsmallest(k,nums)
        #ans[-1]即为第k小的数
        return ans[-1]

结果：

方法二：二分查找

思路：

矩阵虽然在行和列是有序的，但是我们查找的不是某个值，而是第k小的值，因此对矩阵进行二分查找意义不大。

我们对答案值进行二分查找，首先我们想到答案的范围肯定在矩阵的最小值和最大值之间，我们记为minn和maxx，一个为matrix[0] [0]，一个为matrix[n-1] [n-1]。然后我们转换思路，第k小的元素，那么在矩阵中，小于等于它的数的个数为k，或大于k（因为第k小的数在矩阵中可能有重复，导致大于k）。

而对于minn和maxx中的值，随着值的增加，这个值在matrix中小于等于它的数的个数一定是递增的（不严格递增），这就构成了有序的序列，我们就可以对这个序列进行二分查找。注意在二分查找的时候，我们要找的是某个值，这个值在matrix中小于等于它的个数为k，或是第一个大于k的，我们定义一个函数calculate(target)，来计算matrix中小于等于target的个数。

• 如果calculate(mid) >= k，这个mid可能是答案，所以right = mid
• 如果calculate(mid) < k，这个mid不可能是答案，因此left = mid+1，丢弃掉mid

这样最后找到的left=right就是答案，而它对应的calculate值可能等于k也可能大于k（若大于k，是第一个大于k的，对应着答案值在matrix中有多个）。

下面我们来看如何完成calculate函数，我们使用与leetcode240题相似的方法，详细讲解可以看这篇博客
leetcode240. 搜索二维矩阵 II（Python3）

从左下角开始查找，在某个位置x，y时，x表示行序数，y表示列序数。使用cnt计数。

• 如果该位置值小于等于target，那么说明这个位置这一列上面的所有值都小于target，cnt+=x+1。然后向右走一步，y+=1。
• 如果该位置值大于target，说明这个位置的值不符合，需要向上走一步缩小值，x-=1。

当x，y越界，遍历结束，返回cnt。

• 每次calculate的时间复杂度为O(N)，二分查找的时间复杂度为O(log(maxx-minn))，所以总的时间复杂度为O(Nlog(maxx-minn))
• 空间复杂度为O(1)。

代码：

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        minn = matrix[0][0]
        maxx = matrix[n-1][n-1]

        #这个函数的功能是，返回matrix小于等于target的数的个数
        #从左下角开始找。
        def calculate(target):
            #计数
            cnt = 0
            #初始的位置，x为行数，y为列数
            x = n-1
            y = 0
            while 0 <= x < n and 0 <= y < n:
                if matrix[x][y] <= target:
                    cnt += x+1
                    y += 1
                else:
                    x -= 1
            return cnt
        
        #对数值开始二分查找
        left = minn
        right = maxx
        while left < right:
            mid = left + (right-left)//2
            if calculate(mid) >= k:
                right = mid
            #需要向右压缩，因为calculate(mid)<k，则mid可能不是答案
            #若大于k，则mid可能为答案，因为可能存在多个mid这个值
            else:
                left = mid + 1
        return left

结果：