From f02da2c80e9cdb7d597799095b0df2e8e5c7cd19 Mon Sep 17 00:00:00 2001
From: bt3gl <138340846+bt3gl-cryptographer@users.noreply.github.com>
Date: Tue, 8 Aug 2023 16:51:10 -0700
Subject: [PATCH] Update k_smallest_matrix.py

---
 heaps/k_smallest_matrix.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/heaps/k_smallest_matrix.py b/heaps/k_smallest_matrix.py
index 68edf0a..79789cf 100644
--- a/heaps/k_smallest_matrix.py
+++ b/heaps/k_smallest_matrix.py
@@ -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