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

Difficulty
Medium
Tags
二分查找
URL
https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/
Star
给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。 请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。
你必须找到一个内存复杂度优于 O(n2) 的解决方案。
示例 1：
输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
输出：13
解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13
示例 2：
输入：matrix = [[-5]], k = 1
输出：-5
提示：
  • n == matrix.length
  • n == matrix[i].length
  • 1 <= n <= 300
  • 109 <= matrix[i][j] <= 109
  • 题目数据 保证 matrix 中的所有行和列都按 非递减顺序 排列
进阶：
  • 你能否用一个恒定的内存(即 O(1) 内存复杂度)来解决这个问题?
  • 你能在 O(n) 的时间复杂度下解决这个问题吗?这个方法对于面试来说可能太超前了，但是你会发现阅读这篇文章（ this paper ）很有趣。

法 1 二分查找

思路
找出二维矩阵中最小的数 left，最大的数 right，那么第 k 小的数必定在 left ~ right 之间
mid=(left+right) / 2 在二维矩阵中寻找小于等于 mid 的元素个数 count 若这个 count 小于 k，表明第 k 小的数在右半部分且不包含 mid，此时 left=mid+1 若这个 count 大于 等于k，表明第 k 小的数在左半部分且可能包含 mid，此时 , right=mid-1:
所以：寻找的是 左边界， 这个边界右边的数字 temp 都满足 ：整个matrix 中 小于等于 temp 的数字个数 大于等于 k 注意：这里的 left mid right 是数值，不是索引位置。
在收缩过程中，m不一定是 矩阵中的元素.
例如实例 1 ，matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8 , 搜索过程中 m 和 temp (小于等于m 的个数 )的取值分别为：
l, r, m, temp =  1 15 8 2
l, r, m, temp =  9 15 12 6
l, r, m, temp =  13 15 14 8
l, r, m, temp =  13 13 13 8

当 m = 14 的时候， m显然不是矩阵中的元素。
此时 满足 temp <= k，所以此时 r 继续收缩， r = m- 1 = 13,
然后依旧满足，temp <= k , 继续收缩 ，r = m-1 = 12 ,
此时 l > r 了，返回 l, 这也再一次说明了为什么是找左边界
 
题解
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        l, r = matrix[0][0] ,matrix[-1][-1]
        while l <= r:
            m = l + (r - l) // 2
            # 小于等于 m 的个数 
            temp =  self.helper(m, matrix)
            # 小于 k 个：说明肯定在右半部分
            if temp < k:
                l = m + 1
            # 大于等于 k 个：说明在 左半部分，且可能包含mid 
            else:
                r = m - 1

        return l

    def helper(self, m, matrix):
        # 统计  matrix 中小于等于 m 的元素个数
        i, j = 0, len(matrix[0])-1
        count = 0
        while i < len(matrix) and j >= 0:
            # 这里 i += 1 所以先把 i 行的 j + 1 个 元素 加进来
            if matrix[i][j] <= m:
                count += j + 1
                i += 1
            else:
                j -= 1

        return count

### helper函数的另外一种写法
def count(self, matrix, cur_val):
        # 统计每一行中 小于等于 cur_val 的个数
        res = 0
        j = len(matrix[0])-1
        for i in range(0, len(matrix)):
            while j >= 0 and matrix[i][j] > cur_val:
                j -= 1
            res += j + 1
        
        return res