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

题目

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

示例：

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

返回 13。

提示：
你可以假设 k 的值永远是有效的，1 ≤ k ≤ n2 。

解题思路

无脑版：
把二维数组合并成一维，然后返回第k个元素。时间复杂度$o(n^2)$

归并排序版：
合并n个有序数组，然后返回第k个元素。归并的部分和合并K个排序链表思路一致

二分查找版：
二分查找是看了题解才想到的。用low, high分别代表当前矩阵中的最大值和最小值，计算出中值mid = (low + high) // 2，然后找矩阵中<= mid的数字有几个。若<= mid的数字不足k个，则最终的答案应该比mid大。若<= mid的数字>=k个，则最终答案应该<=mid

代码

无脑版：

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        last_list = matrix[0]
        for each_list in matrix[1:]:
            index_1, index_2 = 0, 0
            new_merge = []
            while index_1 < len(last_list) and index_2 < len(each_list):
                if last_list[index_1] < each_list[index_2]:
                    new_merge.append(last_list[index_1])
                    index_1 += 1
                else:
                    new_merge.append(each_list[index_2])
                    index_2 += 1
            while index_1 < len(last_list):
                new_merge.append(last_list[index_1])
                index_1 += 1
            while index_2 < len(each_list):
                new_merge.append(each_list[index_2])
                index_2 += 1
            last_list = new_merge
        return last_list[k - 1]

归并版：

class Solution:
    def merge_2list(self, list1: list, list2: list) -> list:
        ret_list = []
        index1, index2 = 0, 0
        while index1 < len(list1) and index2 < len(list2):
            if list1[index1] < list2[index2]:
                ret_list.append(list1[index1])
                index1 += 1
            else:
                ret_list.append(list2[index2])
                index2 += 1
        if index1 != len(list1):
            ret_list += list1[index1:]
        if index2 != len(list2):
            ret_list += list2[index2:]
        return ret_list

    def merge_sort(self, matrix: List[List[int]]) -> list:
        if not matrix:
            return None
        if len(matrix) == 1:
            return matrix[0]
        mid = len(matrix) >> 1
        return self.merge_2list(self.merge_sort(matrix[:mid]), self.merge_sort(matrix[mid:]))

    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        sorted_list = self.merge_sort(matrix)
        return sorted_list[k - 1]

二分查找版：

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        def not_enough(target: int, k: int) -> bool:
            """
            True: # of numbers that are no larger than target < k
            False: # of numbers that are no larger than target >= k
            """
            i, j = 0, n - 1
            res = 0
            while 0 <= i < n and 0 <= j < n:
                if matrix[i][j] <= target:
                    i += 1
                    res += j + 1
                else:
                    j -= 1
            return res < k

        n = len(matrix)
        left, right = matrix[0][0], matrix[n - 1][n - 1]
        while left < right:
            mid = (left + right) >> 1
            if not_enough(mid, k):
                left = mid + 1
            else:
                right = mid
        return (left + right + 1) >> 1