算法练习：有序矩阵中第K小的元【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 ≤ n2 。

链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix

class Solution:
    def kthSmallest(self, matrix, k):
        n = len(matrix)
        i, j = 0, 0
        start = matrix[0][0]
        end = matrix[n-1][n-1]
        
        index = 0 
        while start < end:
            index = 0
            mid = start + (end - start)//2
            i, j = n-1, 0
            while i >=0 and j < n:  # 局部范围比较

                if matrix[i][j] <= mid:  # 一个优化的地方，找出 mid 的 索引，i 可以从 mid 开始
                    index += 1 + i  # 在第 i 行，从左向右比较,记录索引
                    j += 1
                else:
                    i -= 1          # 自下往上比较，mid 变动后，i 要相应变动
            if index < k:           # 根据局部范围比较的结果，变动 mid
                start = mid + 1
            else:
                end = mid
        return start

题解一：

• 首先这有2个子问题：排序查找，返回索引对应的数；这是一个有序二维表，从左上角到右下角增大;

一般是先排序再查找，有序列表可以使用归并排序，但是这里有 n 维，查找较好的方法是二分查找，现在比较好的方法是在排序过程中查找。

• 可用的特殊条件，n*n 维矩阵，行列都是有序的，则左上角的元素因该是最小的，右下角的元素应该是最大的
• 将 matrix 矩阵元素视为 从 matrix[0][0] 到 matrix[n-1][n-1] 的有序列表； mid = matrix[n-1][n-1]//2,最好为 mid = matrix[0][0] + (matrix[n-1][n-1] - matrix[0][0]),传统情况能避免索引出界。
• 在比较值得过程中，计算出对应值的 index，如果记录的 index 小于 k，就在另一半中查找；最后就返回记录的 index
• 在遍历的时候，还是要做两重，一是将元素大小二分查找范围，缩减范围；而是在缩减的范围内寻找索引的值。

题解二：

刚刚也提到使用归并排序，类似可考虑堆、优先序列

问题可转换为在 N 个有序列表找第 K 小的元素。

如果2个有序列表，找到第 K 个，只需2个指针分别遍历 2 个列表，然后比较大小，找到第 k 个。
假如有N个指针指向每一行，每次取最小，然后将那一行的指针往后偏移，找到第 K 个。每次 N 个取最小，时间复杂度需要 O(N)。

• 堆或优先队列：维护一个 N 大小的最小堆，每次 pop 最小的，只需要 logN。当 k<N 时，只需遍历前k行，只需维护 k 大小的堆。
  时间复杂度：有序列表建堆是O(min(k,N)) + k*O(log(min(k,N)))，空间复杂度：O(min(k,N))

参考：https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix/solution/on-python-by-jerry4free/

from heapq import *

class Solution(object):
    def kthSmallest(self, matrix, k):
        minheap = []
        n = len(matrix)
        for i in range(min(k, n)):  # 当k小于n时，只需遍历前k行
            heappush(minheap, (matrix[i][0], i, 0))  # 堆里的每个元素是（matrix[i][j], i, j）
            
        counter = 0
        x, i, j = 0, 0, 0
        while counter < k:
            counter += 1
            x, i, j = heappop(minheap)
            if j < n-1:
                heappush(minheap, (matrix[i][j+1], i, j+1))  # 向堆里加入该元素所在行的下一个元素
        return x