Kth Smallest Element in a Sorted Matrix

Problem: Kth Smallest Element in a Sorted Matrix

We can add all rows to a heap sorted by the first element of the row. Everytime we want to get the smallest element, we can pop a row from the heap, pop the first element of the row, and push it back. But if the row is very long, the heap may easily exceeds memory limit. To avoid this, we can only store the first element of the row, the index of the row, and index of the current storing element in the row.

Code in Python:

from heapq import heappush, heappop, heapreplace, heapify
class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        h = [(row[0], row, 1) for row in matrix]
        heapify(h)

        # Since we want to find kth, we pop the first k elements 
        for _ in xrange(k - 1):
            v, r, i = h[0]
            if i < len(r):
                heapreplace(h, (r[i], r, i + 1))
            else:
                heappop(h)

        return h[0][0]