master-algorithms-py/heaps/k_smallest_matrix.py
2023-08-08 16:51:10 -07:00

24 lines
503 B
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# author: bt3gl
def kth_smallest(matrix, k) -> int:
min_heap = []
for row in range(min(k, len(matrix))):
min_heap.append((matrix[row][0], row, 0))
heapq.heapify(min_heap)
while k:
element, row, col = heapq.heappop(min_heap)
if col < len(matrix) - 1:
heapq.heappush(min_heap, (matrix[row][cow + 1], row, col + 1))
k -= 1
return element