LeetCode: 378. Kth Smallest Element in a Sorted Matrix

Coming up with O(1) memory solution was very difficult for me.

I refered to this post's last solution.

https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/discuss/1322101/C%2B%2BJavaPython-MaxHeap-MinHeap-Binary-Search-Picture-Explain-Clean-and-Concise

[C++/Java/Python] MaxHeap, MinHeap, Binary Search - Picture Explain - Clean & Concise - LeetCode Discuss

 

One thing you should remember about sorted matrix(both for rows and columns) is that it has a very efficient algorithm to find an element or to count the number of elements smaller/bigger than a target within O(N+M) time complexity and O(1) memory complexity. 

This problem also lies down on the above algorithm.

 

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        m = len(matrix)
        def count(mid):
            cnt = 0
            y, x = 0, m - 1
            while y <= m - 1 and x >= 0:
                if matrix[y][x] <= mid:
                    cnt += x + 1
                    y += 1
                else:
                    x -= 1
            return cnt
        
        left, right = matrix[0][0], matrix[m-1][m-1]
        ret = matrix[m-1][m-1]
        while left <= right:
            mid = (left + right) // 2
            cnt = count(mid)
            if cnt < k:
                left = mid + 1
            else:
                ret = min(ret, mid)
                right = mid - 1
        return ret