378 - heap

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        h = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(h)
        while k:
            val, i, j = heapq.heappop(h)
            if j < n - 1:
                heapq.heappush(h, (matrix[i][j+1], i, j + 1))
            k -= 1

        return val