Kth Smallest Number in Sorted Matrix

Question

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

Thoughts

1. We can build the heap by add all the elements in the matrix one by one. The time complexity is almost O(NlogN), which is not efficient. Considering that another condition is that all the elements are sorted in both the row and column directions. We can utilize to save time when building the heaps.
2. We don't need to build the complete heap structure for all elements. Actually, we only care about the previous k elements in the heap. In other words, we only need to build the heap with around k times actually.

Code

from heapq import heappush, heappop

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
        heap = []
        row = len(matrix)
        col = len(matrix[0])

        for i in xrange(col):
            item = (matrix[0][i], 0, i)
            heappush(heap, item)

        for i in xrange(k - 1):
            item = heappop(heap)
            row_index = item[1]
            col_index = item[2]
            if row_index + 1 < row:
                new_item = (matrix[row_index + 1][col_index], row_index + 1, col_index)
                heappush(heap, new_item)

        return heap[0][0]

Reference

reference source: https://lefttree.gitbooks.io/leetcode/content/dataStructure/heap/kthSmallestInMatrix.html