378.Kth Smallest Element in a Sorted Matrix

Tags: [heap]

Link: https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/\#/description

Given anxnmatrix where each of the rows and columns are sorted in ascending order, find the kth smallest element in the matrix.

Note that it is the kth smallest element in the sorted order, not the kth distinct element.

Example:

matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8,

return 13.

Note:
You may assume k is always valid, 1 ≤ k ≤ n2.

Solution: Heap

import heapq


class Solution(object):

    class HeapNode(object):
        def __init__(self, row, col, val):
            self.row = row
            self.col = col
            self.val = val

        def __cmp__(self, other):
            return self.val - other.val


    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        if not matrix or len(matrix) != len(matrix[0]) or not (1 <= k <= len(matrix) ** 2):
            raise ValueError('The input is invalid')

        num_of_rows = len(matrix)
        num_of_cols = len(matrix[0])

        # load the first row into min_heap
        min_heap = []
        for col in xrange(num_of_cols):
            min_heap.append(self.HeapNode(0, col, matrix[0][col]))

        for _ in xrange(k - 1):
            curr = heapq.heappop(min_heap)
            if curr.row + 1 < num_of_rows:
                heapq.heappush(min_heap, self.HeapNode(curr.row + 1, curr.col, matrix[curr.row + 1][curr.col]))

        return heapq.heappop(min_heap).val

Revelation:

• We know the first row elements are the smallest elements in the matrix, the following elements can be equal to the first row elements, but they cannot be smaller than the first row elements. So we can just use the first row elements to build a min heap. Then do heap pop. And we know the elements in each col are also sorted and also we have already used the row elements, so after we do heap pop, we use the col elements to supply the min heap.
• We can also first load the first col elements, then use the row elements to supply the min heap.
• Since the first row elements have been sorted, so when we load them into min heap, we don't need do heap push, we just use list append.

Note:

• Time complexity = O(n + klgn), 1<= k <= n**2, so the time complexity = O(klgn), n is the number of rows or number of cols.
• In Python, we can define the inner class, but when we use them, we should access them via 'self' keyword.
• In Python, we can override the '__cmp__' function.