378.-kth-smallest-element-in-a-sorted-matrix.md

378. Kth Smallest Element in a Sorted Matrix

• Medium

• 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.

  Note that it is the kth smallest element in the sorted order, not the kth distinct element.

  You must find a solution with a memory complexity better than O(n2).

Analysis

Since we know the next largest value must be to the right of the values we have already seen or is the head of the highest row which we haven checked, we can keep a record of where we have seen in each row and thus we would know where to find the next value.

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n=len(matrix)
        index=[-1]*n
        index[0]=0
        while k>0:
            mini=float('inf')
            minind=0
            for i in range(n):
                if index[i]<n and index[i]!=-1:
                    if matrix[i][index[i]]<mini:
                        mini=matrix[i][index[i]]
                        minind=i
            index[minind]+=1
            if minind<n-1 and index[minind+1]==-1:
                index[minind+1]=0
            k-=1
        return mini

Below is a more advanced solution based on binary search. The range for search is from the smallest number in the matrix [0][0] to the largest one in the matrix[n-1][n-1]. When we search, the thing we are looking for is the smallest value where there are k-1 vaues in the matrix that are not larger than it.

class Solution:  
    def kthSmallest(self, matrix, k):
        m, n = len(matrix), len(matrix[0])  # For general, the matrix need not be a square

        def countLessOrEqual(x):
            cnt = 0
            c = n - 1  # start with the rightmost column
            for r in range(m):
                while c >= 0 and matrix[r][c] > x: c -= 1  # decrease column until matrix[r][c] <= x
                cnt += (c + 1)
            return cnt

        left, right = matrix[0][0], matrix[-1][-1]
        ans = -1
        while left <= right:
            mid = (left + right) // 2
            if countLessOrEqual(mid) >= k:
                ans = mid
                right = mid - 1  # try to looking for a smaller value in the left side
            else:
                left = mid + 1  # try to looking for a bigger value in the right side

        return ans