Kth smallest elment in sorted Matrix

Question

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

Example Given k = 4 and a matrix:

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

return 5

Challenge O(k log n), n is the maximal number in width and height.

Thoughts

refer

1. 利用heap,先对第一行所有元素加入heap,每个元素下面同一列的元素必然比他们大
2. 重复K-1次下面的过程
   • 取现在的root
   • 将root下面的元素加入heap

Solution

from heapq import *
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):
        # write your code here
        if not matrix:
            return 0

        heap = []
        row = len(matrix)
        col = len(matrix[0])

        for i in range(col):
            heappush(heap, (matrix[0][i], 0, i))

        for j in range(k-1):
            curr = heappop(heap)
            x = curr[1]
            y = curr[2]
            if x + 1 < row:
                heappush(heap, (matrix[x+1][y], x+1, y))

        return heap[0][0]