ctci/Chapter 10 Sorting and sear.../10.9. sorted_matrix_search.md

1.1 KiB

10.9. Sorted matrix search

Given an M x N matrix in which each row and each column is sorted in ascending order, write a method to find an element

example

A = [[1,2,4,5,7],
     [2,3,5,6,8],
     [3,6,7,8,10],
     [4,7,8,9,11]]

num = 3

First idea

If A[0][3] < num and A[3][0] < num, if A[3][3] < num, return -1. Otherwise it must be in the matrix.

If A[0][3] => num or A[3][0] => num, Iterate through the last column, find first row where A[row][3] > num. Then iterate that row from the back. If not found, iterate through the last row, do the same. This should take O(N+M) in the worst case.

def search_sorted_matrix(M, N, A, num):
    if num < A[0][0] or num > A[M][N]:
        return 0
    if num == A[0][0] or num == A[M][N]:
        return 1
    if num < A[0][N]:
        initidx = 0
        maxidx = N
        while True:
            i = maxidx // 2
            if A[0][i] == num:
                return 1
            elif A[0][i] > num:
                maxidx = i
            elif A[0][i] < num:
                initidx = i

            if A[0][initidx] < num and A[0][initidx+1] > num:
                break