Leetcode 378: Kth Smallest Element in a Sorted Matrix

Given a n x n matrix 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 ≤ n².

Code

from heapq import *

class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        # each element in the heap: (val, x coord, y coord)
        h = []

        for i in xrange(min(len(matrix[0]), k)):
            heappush(h, (matrix[0][i], 0, i))
        
        for i in xrange(k-1):
            val, x, y = heappop(h)
            if x < len(matrix)-1:
                heappush(h, (matrix[x+1][y], x+1, y)) 

        return h[0][0]

Idea

Maintain a min-heap with k element, initialized by the elements of the first row. Since it is a min-heap, and note the property that rows and columns are already sorted in ascending order, the heap root after popping k-1 times is the k-th smallest element of the whole matrix. When popping the heap, we also need to push necessary matrix elements into the heap.

You can imagine that when you traverse the matrix and maintain the heap, actually you are swiping through a sector area of the matrix, starting from the first row:

sector

Space complexity is O(K) (heap size)

Time complexity is O(KlogK) (every heap operation takes O(logK))

Reference: https://discuss.leetcode.com/topic/52948/share-my-thoughts-and-clean-java-code

Pythonically brute-force:

Code

from heaps import *
class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        return list(merge(*matrix))[k-1]

Note that each argument sent to heapq.merge should already be sorted from smallest to largest. (https://docs.python.org/2/library/heapq.html)

Binary Search

Code

class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        n = len(matrix)
        left, right = matrix[0][0], matrix[n-1][n-1]
        while left < right:
            mid = left + (right - left)/2
            count = 0
            j = n-1
            for i in xrange(n):
                while j >=0 and matrix[i][j] > mid:
                    j -= 1
                count += j + 1
            if count >= k:
                right = mid
            else:
                left = mid + 1
        
        return left

Idea:

We can eventually find the k-th smallest element by shrinking the search range in binary search. Binary search is feasible for this problem since left, right,and mid in binary search are integers and we know that matrix elements are integers.

The algorithm takes O(nlogN) time (N is the range of matrix[0][0] ~ matrix[n-1][n-1]) and O(1) space.

Time complexity analysis: the outer loop (line 10-21) executes at most O(logN) times. The inner for loop (line 14-17) executes at most O(n) times.

Reference: https://discuss.leetcode.com/topic/52865/my-solution-using-binary-search-in-c/16

Also a much more complicated algorithm taking O(N) time: https://discuss.leetcode.com/topic/53126/o-n-from-paper-yes-o-rows

Leetcode 378: Kth Smallest Element in a Sorted Matrix

Leave a comment

Cancel reply