Kth Smallest Element in a Sorted Matrix

Given an n x n matrix where each of the rows and columns are sorted in ascending order, return 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 1:

Input: matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
Output: 13
Explanation: The elements in the matrix are [1,5,9,10,11,12,13,13,15], and the 8th smallest number is 13

Example 2:

Input: matrix = [[-5]], k = 1
Output: -5

Constraints:

• n == matrix.length

Approach 1:

Iterate through entire matrix and see whether the current number “ num” is the k th smallest , by checking the count of elements less than or equal to “num” is k .

Time :O(n³)

Space :O(1)

Approach 2:

Traverse the matrix , store the elements in to a 1 D array . Sort the array and find the kth element in the sorted array

Time :O(n²)

Space :O(n)

Approach 3:

Using max heap

The space complexity can be further reduced to k by using a k-size max heap.

Travers the array and store elements into a k-size max heap. And if the current element is less than the root of max heap , pop the root and push this element to the heap.

import heapq as h
class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
arr=[]
h.heapify(arr)
n=len(matrix)

for i in range(n):
for j in range(n):
if len(arr)<k:
h.heappush(arr,-1*matrix[i][j])
else:
if(arr<-1*matrix[i][j]):
h.heappush(arr,-1*matrix[i][j])
h.heappop(arr)
return arr*-1

Time :O(n²)

Space :O(k)

Approach 4:

Binary Search

This is much more optimized solution in terms of time & space.

As the rows and columns are sorted seperately. We are pretty sure that the smallest element of the matrix will be at , matrix and largest element at the matrix[n-1][n-1].

That is the ranges of numbers falling under the matrix will be from these ranges (matrix,matrix[n-1][n-1]) , inclusive.

So essentially our idea is use a binary search for these ranges, and for current mid , we take count of numbers present in the matrix , less than or equal to this mid value. If that count is less than k , then the the start need to be shifted to mid+1 . Else if the count is greater than k , end is shifted to mid -1 .

The point to note here is , if this count is equal to k , the current mid , need not necessarily be the answer. We should check if there is possibility for mid -1 to be the answer . So in this scenario also end is shifted to mid-1, by storing the previous mid value to a variable “ans”.

Finally return the ans variable

Also note that , if we take the count of elements less than mid value by iterating the entire array, we will have time complexity of O(n²). So some optimization is need to be brought to the counting technique too. Our approach is taking a counting complexity of O(n). This is possible since the column and rows are sorted seperately. Here as we iterate through each array and we store number of elements less than k in that row , to a varaiable “total” by moving the column pointer till the value is less than or equal to k. Initially column pointer j is at the last column.

class Solution:
def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
n=len(matrix)
low,high=matrix,matrix[n-1][n-1]
while(low<=high):
mid=(low+high)//2
cnt= self.counter(matrix,mid)
if cnt<k:
low=mid+1
else:
ans=mid
high=mid-1
return ans

def counter(self,A,k):
total=0
i,j,n=0,len(A)-1,len(A)
while(i<len(A)):
if(A[i]>k):
break
while(k<A[i][j]):
j-=1
total=total+j+1
i+=1