Kth smallest element in a sorted matrix

'''
Logic:
We can declare a min heap and put first elements from all the columns along with cordinates.
min heap will sort the elements based on the first component of each element.
pop top element from min heap and put the next element from that column only.
this way we will get elements one by one and return the kth element.

This method is better when range is big but we are allowed to have O(n) space.
'''

from typing import List
import heapq

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)

        # min heap: (value, row, col)
        heap = []

        # push first row elements (one from each column)
        for col in range(n):
            heapq.heappush(heap, (matrix[0][col], 0, col))

        count = 0

        while heap:
            val, row, col = heapq.heappop(heap)
            count += 1

            if count == k:
                return val

            # push next element from same column
            if row + 1 < n:
                heapq.heappush(heap, (matrix[row + 1][col], row + 1, col))

        return -1


#M-2
'''
This is better when range is small but no extra space allowed.The approach is same as median in the row wise sorted matrix.

mid is not necessarily an element from the matrix, but it represents a value in the search space. By checking how many elements are ≤ mid and shrinking the range accordingly, the binary search converges to the smallest value for which count ≥ k, which is exactly the k-th smallest element in the matrix.

bisect_right(row, mid) returns the index where mid would be inserted on the right, which is equal to the number of elements ≤ mid in that row. For example, if row = [0, 1, 2] and mid = 5, then bisect_right(row, 5) = 3, meaning all 3 elements are ≤ mid.

T:C O(Nlog(N)*log(M)) m=range/space
log(M) is for binary search as we are searching on a space of max-min of the matrix
Nlog(N) is for each mid, we are doing binary search on each row ie. log(N) for each row and N times for N rows so N*log(N)
'''
import math
from bisect import bisect_right
def fun(matrix,target): #This tells how many elements are smaller or equal than the current elemnet
    ans=0
    for row in matrix:
        ans+=bisect_right(row,target)
    return ans
class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        l = matrix[0][0]
        r = matrix[-1][-1]

        while l<=r:
            mid= (l+r)//2
            count=fun(matrix,mid)
            if count>=k:
                r=mid-1
            else:
                l=mid+1
        return l #Mid was just a middle element, not an element from the matrix but you kept reducing when count is more or greater so it came down to the right element of the matrix.