Определение строки с наибольшим количеством единиц в упорядоченной бинарной матрице

Условие

Дана квадратная матрица размера N × N, заполненная нулями и единицами. В каждой её строке после первой встреченной единицы все последующие элементы также являются единицами. Требуется за минимальное время найти индекс строки, содержащей максимальное количество единиц.

Пример матрицы:

arr = [
    [0, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 1],
    [0, 0, 0, 0, 1, 1],
    [0, 0, 0, 1, 1, 1],
    [0, 1, 1, 1, 1, 1],
    [0, 0, 0, 0, 0, 1],
    [0, 0, 0, 1, 1, 1]
]

Ответ

Двигайтесь от правого верхнего угла влево по единицам и вниз по строкам. В примере максимум у строки 4 при нумерации с нуля. Для N×N требуется O(N) шагов.

Пример фактически прямоугольный, 7×6, что не мешает алгоритму:

def row_with_most_ones(matrix):
    if not matrix or not matrix[0]:
        return -1
    column = len(matrix[0]) - 1
    best = 0
    for row_index, row in enumerate(matrix):
        while column >= 0 and row[column] == 1:
            best = row_index
            column -= 1
    return best

Предполагаются равные длины строк и порядок 0…01…1. При равенстве максимумов выбирается первая строка, включая матрицу из одних нулей. В общем случае сложность O(строки + столбцы), память O(1).