search_2d_matrix

Solution

	
	"""
https://leetcode.com/problems/search-a-2d-matrix/description/
You are given an m x n integer matrix matrix with the following two properties:

    Each row is sorted in non-decreasing order.
    The first integer of each row is greater than the last integer of the previous row.

Given an integer target, return true if target is in matrix or false otherwise.

You must write a solution in O(log(m * n)) time complexity.

"""

from typing import List


def search(arr: List[List[int]], target: int) -> bool:
    # rows, cols = len(arr), len(arr[0])
    top, bottom = 0, len(arr) - 1
    while top <= bottom:
        row = top + ((bottom - top) // 2)
        if target < arr[row][0]:
            bottom -= 1
        elif target > arr[row][-1]:
            top += 1
        else:
            break
    if not (top <= bottom):
        return False
    left, right = 0, len(arr[row])
    row = top + ((bottom - top) // 2)
    while left <= right:
        mid = left + ((right - left) // 2)
        if target < arr[row][mid]:
            right = mid - 1
        elif target > arr[row][mid]:
            left = mid + 1
        else:
            return True
    return False

	

Tests

	
	import pytest

from publicmatt.leet.problems.search_2d_matrix import search


@pytest.fixture
def matrix():
    out = [[1, 3, 5, 7], [10, 11, 16, 20], [23, 30, 34, 60]]
    return out


def test_exists(matrix):
    assert search(matrix, 5)
    assert search(matrix, 3)
    assert search(matrix, 10)
    assert search(matrix, 60)
    assert search(matrix, 11)


def test_not_exists_between_rows(matrix):
    assert not search(matrix, 8)


def test_simple_not_exists_same_row(matrix):
    assert not search(matrix, 13)
    assert not search(matrix, 14)
    assert not search(matrix, 18)
    assert not search(matrix, 28)


def test_out_of_bounds(matrix):
    assert not search(matrix, -1)
    assert not search(matrix, 500)