# \[378]\[中等]\[堆]\[二分] 有序矩阵中第K小的元素

## 题目描述

[378. 有序矩阵中第K小的元素](https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix/)

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

示例：

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

返回 13。
```

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

## 解题思路

### 归并排序 + 堆

由于每一行都是排好序的数组, 可以使用**归并排序**的思想, 类似于链表的操作, 从每行的第一个数字开始, 结合**最小堆**.

* 首先将所有行的第一个数字入堆, 每行的第一个数字肯定是这一行最小的, 因此全局最小的一定在其中(这里一定是左上角的值为最小值)
* 此时堆顶的数字为最小值, 将其弹出, 然后将其所在行的下一个数字入堆, 如果这个数字已经是这一行最后一个数字了, 则无需找到新入堆数字, 堆大小-1
* 这样重复$$k-1$$次, 此时堆顶的元素就是整个矩阵中第$$k$$小的元素

```python
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        heap = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(heap)

        for i in range(k - 1):
            num, r, c = heapq.heappop(heap)
            if c < n - 1:
                heapq.heappush(heap, (matrix[r][c + 1], r, c + 1))
        return heap[0][0]
```

时间复杂度为$$O(n^2\log{n})$$, 空间复杂度为$$O(n)$$

其实这道题是[23. 合并K个排序链表](https://leetcode-cn.com/problems/merge-k-sorted-lists/)的简化版, 对于链表, 没有要求同列数字随行递增, 也没有要求每行的数字个数相同, 仍然可以使用归并排序+堆的套路解答. 但这同时也说明我们没有用到这些有用信息, 时间复杂度还有提升的空间.

### 二分查找

二分查找的思路参考: [有序矩阵中第K小的元素](https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix/solution/you-xu-ju-zhen-zhong-di-kxiao-de-yuan-su-by-leetco/)

```python
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        left, right = matrix[0][0], matrix[n - 1][n - 1]  # 既然用二分法, 用`left`和`right`命名最合适
        while left < right:
            mid = (left + right) // 2
            if self.is_left(matrix, mid, k, n):
                right = mid
            else:
                """
                这里count加1是因为: 这里判断是在右半边, 而mid等于left时是被判别为左边的, 因此mid肯定不是最终的结果
                """
                left = mid + 1
        return left

    @staticmethod
    def is_left(mat, mid, k, n):
        i, j = n - 1, 0
        count = 0
        while i >= 0 and j < n:
            if mat[i][j] > mid:
                i -= 1
            else:
                count += i + 1
                j += 1
        return count >= k
```

## 相关题目

* [\[23\]\[困难\]\[堆\] 合并K个排序链表](https://blessbingo.gitbook.io/garnet/suan-fa/lian-biao/23-he-bingkge-pai-xu-lian-biao)