master-algorithms-py/heaps/k_smallest_matrix.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# author: bt3gl

# Given an n x n matrix where each of the rows and columns is sorted 
# in ascending order, return the kth smallest element in the matrix.

def kth_smallest(self, matrix, k) -> int:

        n = len(matrix)
        min_heap = []

        for r in range(min(k, n)):

            min_heap.append((matrix[r][0], r, 0))

        heapq.heapify(min_heap)

        while k:

            element, r, c = heapq.heappop(min_heap)

            if c < n - 1:
                heapq.heappush(min_heap, (matrix[r][c + 1], r, c + 1))
            
            k -= 1
        
        return element