Jeremy Fisher Jeremy Fisher - 2 months ago 58x
Python Question

kth smallest element sorted matrix using QuickSelect

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
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)
return self.matrixRecur(splicedMatrix, (mid_row, mid_col + 1), (end_row, end_col), k-lcount)

I pass in tuples to
which contain the
(row, col)
of the endpoints of the interval. So, if I want to search the whole matrix, I pass
(0, 0)
(n, n)
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
. If
is less than the number of elements less than the midpoint (
), then the kth smallest element is within the left interval because there are at most
elements less than the midpoint and

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.


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.


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