Create k_smallest_matrix.py

This commit is contained in:
marina 2023-08-02 14:21:59 -07:00 committed by GitHub
parent 586803c39b
commit 4ce3ef0daa
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -0,0 +1,28 @@
#!/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