LeetCode 378. 有序矩阵中第K小的元素（二分或者归并）

Description

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

示例： matrix = [ [ 1, 5, 9], [10, 11, 13], [12, 13, 15] ], k = 8, 返回 13。 来源：力扣（LeetCode） 链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

Solution 1 二分

题解

class Solution: def kthSmallest(self, matrix: List[List[int]], k: int) -> int: n = len(matrix) # 注意题目中给定的矩阵是不包含重复元素的，单调升序 def check(mid): i, j = n - 1, 0 num = 0 while i >= 0 and j < n: if matrix[i][j] <= mid: num += i + 1 j += 1 else: i -= 1 return num left, right = matrix[0][0], matrix[-1][-1] while left < right: mid = (left + right) // 2 num = check(mid) if num < k: # 排除mid元素 left = mid + 1 else: right = mid return left

Solution 2 归并

同时遍历k个数组，对同时取出的k个值用小根堆优化取出最小值。

class Solution: def kthSmallest(self, matrix: List[List[int]], k: int) -> int: n = len(matrix) pq = [(matrix[i][0], i, 0) for i in range(n)] heapq.heapify(pq) ret = 0 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]