w3resource

Python: Find the kth smallest element in the matrix

Python heap queue algorithm: Exercise-12 with Solution

Given a n x n matrix where each of the rows and columns are sorted in ascending order, write a Python program to find the kth smallest element in the matrix.

Assume k is always valid, 1 ≤ k ≤ n2 .

Sample Solution:

Python Code:

import heapq
class Solution(object):
    def find_Kth_Smallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        m, n = len(matrix), len(matrix[0])
        temp = [[False] * n for _ in range(m)]
        q = [(matrix[0][0], 0, 0)]
        ans = None
        temp[0][0] = True
        for _ in range(k):
            ans, i, j = heapq.heappop(q)
            if i + 1 < m and not temp[i + 1][j]:
                temp[i + 1][j] = True
                heapq.heappush(q, (matrix[i + 1][j], i + 1, j))
            if j + 1 < n and not temp[i][j + 1]:
                temp[i][j + 1] = True
                heapq.heappush(q, (matrix[i][j + 1], i, j + 1))
        return ans

matrix = [
   [0, 5, 9],
   [11, 12, 13],
   [15, 17, 19]   
]
k = 1
s = Solution()
result = s.find_Kth_Smallest(matrix, k)
print("First smallest element:",result)
k = 2
s = Solution()
result = s.find_Kth_Smallest(matrix, k)
print("\nSecond smallest element:",result)
k = 3
s = Solution()
result = s.find_Kth_Smallest(matrix, k)
print("\nThird smallest element:",result)

Sample Output:

First smallest element: 0

Second smallest element: 5

Third smallest element: 9

Flowchart:

Python heap queue algorithm: Find the kth smallest element in the matrix.

Visualize Python code execution:

The following tool visualize what the computer is doing step-by-step as it executes the said program:


Python Code Editor:

Have another way to solve this solution? Contribute your code (and comments) through Disqus.

Previous: Write a Python program to merge multiple sorted inputs into a single sorted iterator (over the sorted values) using Heap queue algorithm.
Next: Write a Python program to find the nth super ugly number from a given prime list of size k using Heap queue algorithm.

What is the difficulty level of this exercise?

Test your Python skills with w3resource's quiz

Python: Tips of the Day

Creates a dictionary with the same keys as the provided dictionary and values generated by running the provided function for each value:

Example:

def tips_map_values(obj, fn):
  ret = {}
  for key in obj.keys():
    ret[key] = fn(obj[key])
  return ret
users = {
  'Owen': { 'user': 'Owen', 'age': 29 },
  'Eddie': { 'user': 'Eddie', 'age': 15 }
}

print(tips_map_values(users, lambda u : u['age'])) # {'Owen': 29, 'Eddie': 15}

Output:

{'Owen': 29, 'Eddie': 15}