Git-Mirrors/master-algorithms-py
1
0
Fork 0
mirror of https://github.com/autistic-symposium/master-algorithms-py.git synced 2025-04-29 12:16:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

Update k_smallest_matrix.py

Browse Source
This commit is contained in:
bt3gl 2023-08-08 16:51:10 -07:00 committed by GitHub
parent 66cf6cc427
commit f02da2c80e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
heaps/k_smallest_matrix.py
View File

@ -2,27 +2,22 @@
# -*- 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)
def kth_smallest(matrix, k) -> int:
min_heap = []
for r in range(min(k, n)):
min_heap.append((matrix[r][0], r, 0))
for row in range(min(k, len(matrix))):
min_heap.append((matrix[row][0], row, 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))
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