Search a 2D Matrix

Problem

Qn - Search a 2D Matrix

Approach

Approach is easy, I guess the test here is on the implementation. We just need to find the correct row using the first element, as the first element will be the smallest number of the row, and larger than the first element of the previous row. Then we just do a binary search on the row since they are sorted.

Complexity

• Time: O(logm * logn)
• Space: O(1)

Code

class Solution:
    def searchMatrix(self, matrix: List[List[int]], target: int) -> bool:
        rows, cols = len(matrix), len(matrix[0])
        
        # find the row to check
        # returns index of the row
        # matrix[i][0] <= target and matrix[i+1][0] > target
        def binary_search_row():
            l, r = 0, rows-1
            while l < r:
                mid = (l+r) // 2
                if matrix[mid][0] <= target:
                    if mid+1 < rows and matrix[mid+1][0] > target: return mid
                    elif mid+1 < rows: l = mid+1
                    else: return mid
                else: r = mid-1
 
            return l
 
        # given a row, find if the target exists in the row
        # returns boolean
        def binary_search_col(row):
            l, r = 0, cols-1
 
            while l < r:
                mid = (l+r) // 2
                if matrix[row][mid] == target: return True
                elif matrix[row][mid] > target: r = mid - 1
                else: l = mid + 1
            
            return matrix[row][l] == target
 
        row = binary_search_row()
        return binary_search_col(row)