Problem statement

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/

Solution

The idea is for number X find number of elements which are less or equal than X, which can be done in O(n), idea is similar to problem 240: Search a 2D Matrix II: we can start with the top right element and move only down and to the left, counting number of elements <X in each row. We start with the smallest and the biggest elements in our table and do binary search, each time asking question: is number of elements < X is more than k? We do binary search and stop when end become equal to beg.

Complexity

Time complexity is O(n * log(A)), where A is difference between maximum and minimum values in our matrix.

Code

class Solution:
    def kthSmallest(self, matrix, k):
        n, beg, end = len(matrix), matrix[0][0], matrix[-1][-1]
        
        def check(m):
            i, j, cnt = 0, n-1, 0
            for i in range(n):
                while j >= 0 and matrix[i][j] > m: j -= 1
                cnt += (j + 1)
            return cnt
         
        while beg < end:
            mid = (beg + end)//2
            if check(mid) < k:
                beg = mid + 1
            else:
                end = mid
                
        return beg

Remark

There is also solution, based on the problem 373: Find K Pairs with Smallest Sums, using heaps with complexity O(k * log k). Depending on matrix it can be bigger or smaller than previous approach.

Finally, there is O(n) time complexity algorithm with the idea of median of medians expanded for 2d matrices, you can find it here https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/85170/O(n)-from-paper.-Yes-O(rows), It is very difficult but I enjoyed to read it.