Подсчёт количества нулей в упорядоченной квадратной матрице

Условие

Реализуйте функцию count_zeros(), принимающую квадратную матрицу, заполненную только символами 0 и 1. В каждой строке и каждом столбце элементы расположены в неубывающем порядке (все нули находятся слева/вверху, а единицы — справа/внизу). Функция должна вернуть общее число нулей, присутствующих во всей матрице.

Пример

matrix = [
    [0, 0, 1, 1],
    [0, 1, 1, 1],
    [1, 1, 1, 1],
    [1, 1, 1, 1]
]
print(count_zeros(matrix))  # вывод: 3

Ответ

Идите от правого верхнего угла: на единице сдвигайтесь влево, на нуле добавляйте количество элементов слева вместе с ним и переходите на следующую строку.

def count_zeros(matrix):
    n = len(matrix)
    if any(len(row) != n for row in matrix):
        raise ValueError('Ожидалась квадратная матрица')
    row, column, total = 0, n - 1, 0
    while row < n and column >= 0:
        if matrix[row][column] == 0:
            total += column + 1
            row += 1
        else:
            column -= 1
    return total

Сортировка по строкам гарантирует нули слева от найденного нуля; сортировка по столбцам позволяет не возвращаться вправо. Выполняется не более 2n перемещений, память O(1). Код рассчитан на числа 0 и 1 из примера; для строковых символов сравнивайте с '0'. Проверка всех значений и сортировки потребовала бы O(n²).