kth smallest element sorted matrix using QuickSelect

Question

I know this has been asked before but this question is about my specific code. I am trying to do a psuedo QuickSelect algorithm where I compare k to the midpoint of a sub interval of a sorted matrix.

I keep getting a timeout error.

Here is the matrix:

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

Here is my code:

def kthSmallest(self, matrix, k):
    """
    :type matrix: List[List[int]]
    :type k: int
    :rtype: int
    """

    return self.matrixRecur(matrix, (0, 0), (len(matrix) - 1, len(matrix[0]) - 1), k)

def matrixRecur(self, splicedMatrix, left, right, k):
    start_row, start_col = left
    end_row, end_col = right
    mid_row = (start_row + end_row)/2
    mid_col = (start_col + end_col)/2

    i = start_row
    j = start_col
    lcount = 0
    while(not (i == mid_row and j == mid_col)):
        if j < len(splicedMatrix[0]):
            j += 1
        else:
            j = 0
            i += 1
        lcount += 1
    if k == lcount:
       return splicedMatrix[mid_row][mid_col]
    elif k < lcount:
        return self.matrixRecur(splicedMatrix, (start_row, start_col), (mid_row, mid_col), k)
    else:
        return self.matrixRecur(splicedMatrix, (mid_row, mid_col + 1), (end_row, end_col), k-lcount)

I pass in tuples to matrixRecur which contain the (row, col) of the endpoints of the interval. So, if I want to search the whole matrix, I pass (0, 0) and (n, n) . matrixRecur will look at a subinterval, determine the midpoint based on the row col numbers of the endpoints, count the number of elements less than the midpoint, compare it to k . If k is less than the number of elements less than the midpoint ( lcount ), then the kth smallest element is within the left interval because there are at most lcount elements less than the midpoint and k < lcount .

I am running this on an interview question site and the site continues to tell me my code times out. I am not seeing why.

Answer 1

Your above approach will not work. As your matrix is row and column wise sorted. Consider the matrix like below

10, 20, 30, 40
15, 25, 35, 45
24, 29, 37, 48
32, 33, 39, 50

Your approach will fail in this case. You are getting time out because you are traversing the whole 2D matrix. Worst case time complexity would be O(mn) (m and n are respectively number of rows and columns).

You can use a min-heap to solve this problem.

Algorithm:

1. Build a min heap of first row elements. A heap entry also stores row number and column number.

2. for(int i=0;i<k;i++)
       Get minimum element (or root) from min heap.
       Find row number and column number of the minimum element.
       Replace root with the next element from same column and min-heapify the root.

3. Return the last extracted root.

Time complexity : O(n + kLogn)

Source : Kth smallest in 2D matrix