Kth Smallest Number in Sorted Matrix

Posted 2020-07-28 Sheryl Wang

Find the kth smallest number in at row and column sorted matrix.

Given k = 4 and a matrix:

[
  [1 ,5 ,7],
  [3 ,7 ,8],
  [4 ,8 ,9],
]

return 5

这一题是Kth Largest Element in an Array的拓展，主要是在排序数组中我们如何从小到大选择，最终到达第k小的这个元素．比如[0,0]位置的这个元素肯定是最小值，之后，备选的元素就是坐标[1,0],[0,1]的两个点，比较这两个获取第２小的元素．我们选取了[1,0]这个位置的，之后[2,0],[1,1]也称为了候选的元素，所以这个候选集是动态增长的，需要实时选择最小元素．和Super Ugly Number一样，这时heap是一个非常好的选择，一方面可以保存之前已经是备选的元素留着后面继续比较，利用之前元素本身的相对关系，不用每次全部重新排序．堆的大小最大为n,　即每行最多只有一个元素在堆中，所以每次获得堆中最小元素的时间复杂度为O(logn),

所以最终总的时间复杂度为Ｏ(klogn). 另外在二维矩阵中通过一点拓展到相邻顶点总是会遇到重复遍历的情况，所以维护一个visited矩阵，判断节点是否被遍历比较重要．否则重复添加，反复pop，最终找到的点不对．

具体的实现代码如下：

class Solution:
    # @param matrix: a matrix of integers
    # @param k: an integer
    # @return: the kth smallest number in the matrix
    def kthSmallest(self, matrix, k):
        if not matrix or not matrix[0] or len(matrix)*len(matrix[0]) < k:
            return -1
        m = len(matrix)
        n = len(matrix[0])
        from heapq import *
        i = 0
        heap = [(matrix[0][0],(0,0))]
        visited = [[False]*n for j in xrange(m)]
        visited[0][0] = True
        while i < k:
            cur, (x, y)= heappop(heap)
            i += 1
            if x + 1 < m and visited[x + 1][y] == False:
                heappush(heap, (matrix[x + 1][y], (x + 1, y)))
                visited[x + 1][y] = True
            if y + 1 < n and visited[x][y + 1] == False:
                heappush(heap, (matrix[x][y + 1],(x, y + 1))) 
                visited[x][y + 1] = True
        return cur