LC 0378 [M] Kth Smallest Element in a Sorted Matrix - ALawliet/algorithms GitHub Wiki

• binary search using top-left and bot-right values
• check bisected row sum < k

---

Since each of the rows in matrix are already sorted, we can understand the problem as finding the kth smallest element from amongst M sorted rows.

We binary search to find the smallest ans in range [minOfMatrix..maxOfMatrix] such that countLessOrEqual(ans) >= k, where countLessOrEqual(x) is the number of elements less than or equal to x.

---

# T: O(nlogn), S: O(1)
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        l, r = matrix[0][0], matrix[-1][-1] # top-left, bot-right
        while l < r:
            m = (l + r) // 2
            bisect_row_sum = sum(bisect_right(row, m) for row in matrix)
            # print(m, bisect_row_sum)
            if not bisect_row_sum >= k: 
                l = m + 1
            else:
                r = m
        return l

# T: O(n^2*logn), S: O(n)
class Solution(object):
    def kthSmallest(self, matrix, k):
        heap = [(row[0], i, 0) for i, row in enumerate(matrix)]
        heapq.heapify(heap)
        ret = 0
        for _ in range(k):
            ret, i, j = heapq.heappop(heap)
            if j+1 < len(matrix[0]):
                heapq.heappush(heap, (matrix[i][j+1], i, j+1))
        return ret