Kth Smallest Element in a Sorted Matrix in Python

Suppose we have a n x n matrix where each of the rows and columns are sorted in increasing order, we have to find the kth smallest element in the matrix. Note that it is the kth smallest element in the sorted order, not the kth unique element. So if the input is like [[1,5,9],[10,11,13],[12,13,15]], if k = 8, then the output will be 13.

To solve this, we will follow these steps −

• define one method called checkVal() and the arguments are matrix and value
• i := 0, j := length of matrix[0] – 1, counter := 0
• while i < length of matrix and j >= 0
  • if matrix[i, j] > value, then decrease j by 1, otherwise counter := counter + j + 1, increase i by 1
• return counter
• the main method will be like −
• n := row of matrix, high := bottom right corner element, low := top left corner element
• while low <= high, do
  • mid = low + (high – low)/2
  • count := checkVal(matrix, mid)
  • if count < k, then low := mid + 1, otherwise high := mid – 1
• return low

Let us see the following implementation to get better understanding −

Example

 Live Demo

class Solution(object):
   def kthSmallest(self, matrix, k):
      """
      :type matrix: List[List[int]]
      :type k: int
      :rtype: int
      """
      n = len(matrix)
      high = matrix[n-1][n-1]
      low = matrix[0][0]
      while low<=high:
         mid = low + (high - low) /2
         count = self.check_value(matrix,mid)
         if count< k:
            low = mid+1
         else :
            high = mid-1
      return int(low)
   def check_value(self, matrix, value):
      i = 0
      j = len(matrix[0])-1
      counter = 0
      while(i<len(matrix) and j >=0):
         if matrix[i][j] > value:
            j-=1
         else:
            counter+=j+1
            i+=1
      return counter
matrix = [[1,5,9],[10,11,13],[12,13,15]]
ob = Solution()
print(ob.kthSmallest(matrix, 8))

Input

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

Output

13