python 提取矩阵每个元素

378. 有序矩阵中第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 。

解题思路

今天没有再遇到很熟悉的题目，不过乍一看还是很眼熟，很像剑指offer的第一题。区别还是有的，主要区别是这道题要难一些。。。因为和剑指offer第一题很像，所以第一反应就是锯齿状对角线来这个二维遍历数组，但是没想明白这个思路该怎么继续下去，遂作罢；然后看到了部分有序，很快又反映到二分上，但是仍没想明白该怎么继续，又作罢（当然最后看官方题解时用的是二分+锯齿遍历，Orz，后面会有说到）。那么既然是部分有序，所以就又想到归并，利用归并将这个部分有序的二维数组打平，变成一维的有序数组，然后第K小就是数组中[k-1]位置的元素。因为不用完全打平，只要到[k-1]位置上就可以了，所以cost要比纯暴力全打平稍好一点点，不过也就好一点点。

具体来说就是先取二维数组matrix的前两行做归并成一个有序数组，然后再用这个数组和第三行继续归并，一直循环，有点像spark里rdd.reduce(lambda a,b: a+b)这种操作。截止条件是当下一行第一个元素大于已排序数组中[k-1]位置元素的时候，因为此时归并插入元素的位置都在第K小元素之后，前面的顺序不会再变了，就可以停止了。时间复杂度大约是

O(klogn)还大一些吧，对不起，我时间复杂度这块确实很不好。如果一直没有触发截止条件，则需要将整个二维数组完全打平，k的最坏值为n^2，所以时间复杂度最坏O(n^2logn)。其实单就归并思路来说，这种方法也不是最好的，在归并步骤上还可以改进。在这个方法我是逐行依次归并的，其实应该可以多行两两归并，官方题解中说要用小根堆什么的，还给出了参考例题：leetccde 23. 合并K个排序链表。我看官方题解的代码里似乎调用了现有的方法，实际面试可能不让这么做，所以也没细看。具体如何更好地归并，还是等遇到题目时再说吧，毕竟还要看这道题的官方题解呢，怕一下太多吃不透，谁让我小白呢。那其实多行两两归并的时候复杂度应该就稳定在O(klogn)，而不像这个方法，要比这个时间复杂度大一些。

题解代码：

def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        # 归并排序
        def guibing(s1,s2):  
        # 我知道函数内套函数很蠢，但是我觉得面试时候面试官应该是会让手撕归并的，况且也不难是吧
            res = []
            i = 0
            j = 0
            while i < len(s1) and j < len(s2):
                if s1[i] < s2[j]:
                    res.append(s1[i])
                    i += 1
                else:
                    res.append(s2[j])
                    j += 1
            if i == len(s1):
                res +=s2[j:] 
            else:
                res += s1[i:]
            return res
        
        tmp = matrix[0]
        for i in range(1,len(matrix)):
            nums = matrix[i]
            tmp = guibing(tmp,nums)  # 依次逐行归并
            p = len(tmp)
            if i < len(matrix)-1 and p >= k:  
            '''保证不是最后一行防止超出index的范围；
                同时保证已归并的有序数组长度大于k才会判断是否到达截止条件'''
                if matrix[i+1][0] > tmp[k-1]:  # 到达截止条件
                    return tmp[k-1]
        return tmp[k-1]

运行结果：

不过这方法真的很烂啊，效率可太糟糕了。不过也正常，因为这个方法只利用了每行元素递增这一个性质，每列也递增这个性质并没有用到。那么官方题解就都用到了，所以效果就很好。

官方思路

前面有说官方思路是二分+锯齿搜索。大体来说和剑指offer第一题很相似，就是从右上或者左下开始锯齿状遍历数组内的元素，与所pick的数进行比较，要求数组内左边的数小于pick的数，右边的数大于pick的数，并且上边的数小于pick的数，下边的数大于pick的数。将这一系列临界点连起来就是锯齿状的边界。在遍历中每一次比较的时候，如果pick的数大于当前遍历到的数组内的数，则这个数组内的数，它的左边和上边的所有数都小于pick的数，那么所有左边和上边的数的个数如果小于k，说明我们所pick的数并不合适，不是第k小的数。那么这个pick的数怎么找呢，就要用到二分了。二分的左边界是matrix[0][0]，右边界是matrix[-1][-1]，因为所有数的范围都不会超过这两个边界。然后直接用二分确定pick的数mid = (left + right) // 2。这样做下来时间复杂度在

O(nlog(right−left))，因为二分查找进行次数为O(log(right−left))，每次操作锯齿状遍历的时间复杂度为 O(n)。

具体解释起来有点麻烦，主要是很难说清楚，官方题解里解释得非常清楚，还有图，所以我就不做搬运工了，直接贴上链接：官方题解，记得看的是方法三哈。代码是我看完题解自己写的，不过对照下来似乎一模一样，细节的解释部分我都放在注释里了，应该还算清楚。

def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)

        def checknum(mid):
            # check数组中比mid小的元素是否多于k个
            i = n-1  # 设定初始位置，从左下角开始
            j = 0
            counter = 0  # 记录其左上部分有几个元素，即比mid小的有几个元素
            while i >= 0 and j < n:
                if matrix[i][j] <= mid:  # 未到边界且mid较大，按列记，先将这一列的元素计入
                    counter += i+1  # +1是因为i是index，实际个数要在index上+1
                    j += 1  # 右移
                else:
                    i -= 1  # 上移
            return counter >= k  # 比mid小的元素个数与k比较

        left = matrix[0][0]
        right = matrix[-1][-1]  # 定义左右边界
        while left < right:  # 循环内进行二分
            mid = (left + right) // 2
            if checknum(mid):
                right = mid
            else:
                left = mid+1  # 注意每次更新的边界，别把mid漏了
        return left

运行结果：

以上就是我，一个正儿八经的小白（大神们通过看代码应该也感觉出来了），对这道题的理解，欢迎诸位指正讨论，感谢阅读。