Set Matrix Zeroes: обнулить строки и столбцы за O(1) памяти

Коротко

Задача выглядит как упражнение на два вложенных цикла, и первое решение действительно пишется за минуту. Только оно неверное. Вся сложность Set Matrix Zeroes в том, чтобы не испортить данные, которые ещё не прочитаны. Это тот же вопрос, что и в Move Zeroes, только в двух измерениях.

Условие

Дана матрица m × n из целых чисел. Если элемент равен нулю, нужно обнулить всю его строку и весь его столбец. Менять матрицу надо на месте.

Вход:            Выход:
1 1 1            1 0 1
1 0 1     →      0 0 0
1 1 1            1 0 1

Ещё пример, где нули стоят на краю:

Вход:            Выход:
0 1 2 0          0 0 0 0
3 4 5 2   →      0 4 5 0
1 3 1 5          0 3 1 0

Это задача 73 на LeetCode, уровень Medium. В follow-up прямо сказано: решение с $O(mn)$ памяти плохое, с $O(m + n)$ лучше, а можно ли за константу? Ровно эту лестницу мы и пройдём.

Решение, которое ломается

Самая очевидная мысль: идём по матрице, встретили ноль — обнуляем строку и столбец.

def set_zeroes_naive(matrix):
    m, n = len(matrix), len(matrix[0])
    for i in range(m):
        for j in range(n):
            if matrix[i][j] == 0:
                for k in range(n):
                    matrix[i][k] = 0
                for k in range(m):
                    matrix[k][j] = 0
    return matrix


print(set_zeroes_naive([[1, 0, 1], [1, 1, 1], [1, 1, 1]]))
# [[0, 0, 0], [0, 0, 0], [0, 0, 0]]

Ожидали получить нули в первой строке и втором столбце, а получили полностью нулевую матрицу. Ноль в [0][1] обнулил всю первую строку, включая [0][2]. Дальше цикл дошёл до [0][2], увидел там ноль и решил, что он был в исходной матрице. Обнулил третий столбец, и так по цепочке.

Цикл не отличает исходные нули от тех, что поставил сам. Ответ зависит только от исходных нулей, а мы теряем эту информацию в момент первой записи. Отсюда правило на всю статью: сначала прочитать, где нули, потом писать. Все решения ниже — это разные способы сохранить прочитанное.

Шаг 1: копия матрицы, O(mn) памяти

Самый прямой способ не путать старые нули с новыми — смотреть на копию.

def set_zeroes_copy(matrix):
    m, n = len(matrix), len(matrix[0])
    original = [row[:] for row in matrix]
    for i in range(m):
        for j in range(n):
            if original[i][j] == 0:
                for k in range(n):
                    matrix[i][k] = 0
                for k in range(m):
                    matrix[k][j] = 0
    return matrix

Это правильно, но плохо сразу по двум осям. Память $O(mn)$ на копию. Время тоже хуже, чем кажется: на каждый ноль тратим $m + n$ операций, и если нулей много, выходит $O(mn \cdot (m + n))$. Матрица из одних нулей размером 1000 × 1000 даст порядка двух миллиардов присваиваний.

Копировать строку нужно именно так, row[:] или list(row). Запись original = matrix[:] копирует только внешний список, строки останутся общими, и баг из предыдущего раздела вернётся.

Это решение полезно как эталон для тестов: оно настолько простое, что в нём трудно ошибиться. Все варианты ниже я проверял сравнением с ним на тысячах случайных матриц.

Шаг 2: два множества, O(m + n) памяти

Посмотрим, что на самом деле нужно знать. Копия хранит каждое значение, но ответ зависит от гораздо меньшего: для каждой строки — есть ли в ней ноль, и для каждого столбца — то же самое. Это $m + n$ булевых значений.

def set_zeroes_sets(matrix):
    m, n = len(matrix), len(matrix[0])
    zero_rows, zero_cols = set(), set()

    for i in range(m):
        for j in range(n):
            if matrix[i][j] == 0:
                zero_rows.add(i)
                zero_cols.add(j)

    for i in range(m):
        for j in range(n):
            if i in zero_rows or j in zero_cols:
                matrix[i][j] = 0
    return matrix

Два прохода, оба по $mn$ ячеек, проверка в множестве за $O(1)$ в среднем (почему так, разобрано в статье о хеш-таблицах). Итого $O(mn)$ по времени и $O(m + n)$ по памяти.

Вместо множеств можно взять два списка [False] * m и [False] * n — так даже чуть быстрее и честнее по памяти. Множества я оставил, потому что так нагляднее мысль «запомнили номера строк».

На собеседовании это решение — нормальная остановка. Многие интервьюеры на нём и закончат, если оно написано чисто. Но follow-up про $O(1)$ задают часто, поэтому идём дальше.

Тупиковая ветка: метить «особым значением»

Частая идея на пути к $O(1)$: при встрече нуля ставить в его строку и столбец не ноль, а метку, которой точно нет в матрице. В конце заменить метки на нули.

def set_zeroes_sentinel(matrix):
    MARK = None
    m, n = len(matrix), len(matrix[0])
    for i in range(m):
        for j in range(n):
            if matrix[i][j] == 0:
                for k in range(n):
                    if matrix[i][k] != 0:
                        matrix[i][k] = MARK
                for k in range(m):
                    if matrix[k][j] != 0:
                        matrix[k][j] = MARK
    for i in range(m):
        for j in range(n):
            if matrix[i][j] is MARK:
                matrix[i][j] = 0
    return matrix

Обратите внимание на != 0: исходные нули трогать нельзя, иначе они превратятся в метки и мы их потеряем.

Это работает, но у решения две проблемы. Во-первых, время снова $O(mn \cdot (m + n))$. Во-вторых, «значение, которого точно нет» — допущение. В Python None в списке чисел сработает, но в Java или C++ матрица int[][] не может хранить ничего, кроме int, а LeetCode даёт значения во всём диапазоне от $-2^{31}$ до $2^{31} - 1$. Свободного числа для метки нет. Если на собеседовании предложить -1 или Integer.MIN_VALUE, первым вопросом будет «а если оно встретится во входе?».

Метка нужна, но хранить её надо в таком месте, где её смысл не перепутать с данными.

Шаг 3: маркеры в первой строке и столбце, O(1) памяти

Вернёмся к решению с двумя массивами: строки отмечаются в массиве длины $m$, столбцы — в массиве длины $n$. А в самой матрице уже есть столбец длины $m$ и строка длины $n$: первый столбец и первая строка. Используем их как эти массивы.

• matrix[i][0] == 0 означает «строку i нужно обнулить».
• matrix[0][j] == 0 означает «столбец j нужно обнулить».

Почему так можно? Если в строке i есть ноль, то matrix[i][0] всё равно станет нулём в ответе. Мы записываем туда ноль раньше времени, но значение верное, и ничего ценного не теряем. С первой строкой то же самое.

Остаётся вопрос: а что с самой первой строкой и первым столбцом? Их ячейки теперь заняты маркерами для других, и по ним уже не понять, был ли там исходный ноль. Значит, это надо выяснить до того, как начнём писать маркеры, и сохранить в двух булевых переменных.

def set_zeroes(matrix):
    m, n = len(matrix), len(matrix[0])
    first_row_zero = any(matrix[0][j] == 0 for j in range(n))
    first_col_zero = any(matrix[i][0] == 0 for i in range(m))

    # 1. Помечаем: строку i — в matrix[i][0], столбец j — в matrix[0][j]
    for i in range(1, m):
        for j in range(1, n):
            if matrix[i][j] == 0:
                matrix[i][0] = 0
                matrix[0][j] = 0

    # 2. Обнуляем внутреннюю часть по маркерам
    for i in range(1, m):
        for j in range(1, n):
            if matrix[i][0] == 0 or matrix[0][j] == 0:
                matrix[i][j] = 0

    # 3. Последними — сами маркеры
    if first_row_zero:
        for j in range(n):
            matrix[0][j] = 0
    if first_col_zero:
        for i in range(m):
            matrix[i][0] = 0
    return matrix

Обратите внимание на диапазоны. Шаги 1 и 2 идут с индекса 1, а не с 0: первая строка и первый столбец здесь не данные, а служебная память. Исходные нули, которые в них стояли, уже работают как маркеры без всяких действий с нашей стороны. Ноль в matrix[0][3] сам по себе говорит «обнули столбец 3», и это как раз правда.

Память — две переменные, $O(1)$. Время — несколько проходов по матрице, $O(mn)$.

Разбор на примере

Возьмём матрицу, где ноль стоит в первом столбце:

1 2 3
0 5 6
7 8 0

Сначала флаги. В первой строке нулей нет: first_row_zero = False. В первом столбце есть ноль [1][0]: first_col_zero = True.

Шаг 1. Во внутренней части (строки 1–2, столбцы 1–2) один ноль, [2][2]. Ставим маркеры matrix[2][0] = 0 и matrix[0][2] = 0:

1 2 0
0 5 6
0 8 0

Маркер строки 1 уже стоит: это исходный ноль [1][0].

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

1 2 0
0 0 0
0 0 0

Шаг 3. first_row_zero ложно, первую строку не трогаем. first_col_zero истинно, обнуляем первый столбец:

0 2 0
0 0 0
0 0 0

Проверим по определению. Исходные нули — [1][0] и [2][2]. Значит, нулями должны стать строки 1 и 2 и столбцы 0 и 2. Единственная ячейка вне них — [0][1], там и осталась двойка. Сходится.

Ловушка с первой ячейкой

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

def set_zeroes_broken(matrix):
    m, n = len(matrix), len(matrix[0])
    for i in range(m):
        for j in range(n):
            if matrix[i][j] == 0:
                matrix[i][0] = 0
                matrix[0][j] = 0
    for i in range(1, m):
        for j in range(1, n):
            if matrix[i][0] == 0 or matrix[0][j] == 0:
                matrix[i][j] = 0
    if matrix[0][0] == 0:
        for j in range(n):
            matrix[0][j] = 0
        for i in range(m):
            matrix[i][0] = 0
    return matrix


print(set_zeroes_broken([[1, 2, 3], [0, 5, 6], [7, 8, 0]]))
# [[0, 0, 0], [0, 0, 0], [0, 0, 0]] — двойка пропала

Ячейка matrix[0][0] одновременно маркер строки 0 и маркер столбца 0. Ноль в [1][0] пишет matrix[0][0] = 0 как маркер столбца 0. В конце код читает ту же ячейку и решает, что обнулять надо и первую строку. Два разных факта уместить в один бит нельзя.

Отсюда и нужны две независимые переменные. Минимально достаточно одной: matrix[0][0] честно отдаём, например, под первую строку, а для первого столбца заводим флаг. Этот вариант часто встречается в разборах:

def set_zeroes_one_flag(matrix):
    m, n = len(matrix), len(matrix[0])
    first_col_zero = False
    for i in range(m):
        if matrix[i][0] == 0:
            first_col_zero = True
        for j in range(1, n):
            if matrix[i][j] == 0:
                matrix[i][0] = 0
                matrix[0][j] = 0

    for i in range(m - 1, -1, -1):
        for j in range(n - 1, 0, -1):
            if matrix[i][0] == 0 or matrix[0][j] == 0:
                matrix[i][j] = 0
        if first_col_zero:
            matrix[i][0] = 0
    return matrix

Он короче, но держать в голове больше. Второй проход идёт снизу вверх, чтобы первая строка, где лежат маркеры столбцов, обнулялась последней. Внутренний цикл не заходит в столбец 0, потому что там маркеры строк, а сам столбец обнуляется по флагу после обработки строки. На собеседовании я бы писал версию с двумя флагами: в ней каждый шаг очевиден, а разница в одну переменную ничего не меняет в асимптотике.

Почему порядок шагов важен

Вторая ловушка тише. Шаг 3 должен идти после шага 2. Если сначала обнулить первую строку и столбец, а потом обходить внутреннюю часть, маркеры станут нулями и обнулят лишнее:

Вход: [[0, 1], [1, 1]]

Правильно:          [[0, 0], [0, 1]]
Шаг 3 раньше шага 2: [[0, 0], [0, 0]]

Ноль в [0][0] должен обнулить первую строку и первый столбец, ячейка [1][1] не задета. Но если первая строка уже обнулена, на шаге 2 matrix[0][1] == 0 выглядит как маркер «обнули столбец 1», и единица в углу пропадает. Это тот же баг, что в наивном решении: записали раньше, чем дочитали.

Та же версия на JavaScript

function setZeroes(matrix) {
  const m = matrix.length
  const n = matrix[0].length
  const firstRowZero = matrix[0].some((x) => x === 0)
  const firstColZero = matrix.some((row) => row[0] === 0)

  for (let i = 1; i < m; i++) {
    for (let j = 1; j < n; j++) {
      if (matrix[i][j] === 0) {
        matrix[i][0] = 0
        matrix[0][j] = 0
      }
    }
  }
  for (let i = 1; i < m; i++) {
    for (let j = 1; j < n; j++) {
      if (matrix[i][0] === 0 || matrix[0][j] === 0) matrix[i][j] = 0
    }
  }
  if (firstRowZero) matrix[0].fill(0)
  if (firstColZero) for (let i = 0; i < m; i++) matrix[i][0] = 0
  return matrix
}

console.log(
  JSON.stringify(
    setZeroes([
      [1, 2, 3],
      [0, 5, 6],
      [7, 8, 0],
    ])
  )
)
// [[0,2,0],[0,0,0],[0,0,0]]

Сравнение решений

Быстрее $O(mn)$ не бывает: чтобы найти все нули, каждую ячейку надо прочитать хотя бы раз. Так что спор идёт только о памяти. Как считать такие оценки, если нужно освежить, — в разборе O-нотации.

Небольшая оговорка про «$O(1)$». Решение с маркерами экономит память, но портит первую строку и столбец на время работы. Если матрицу параллельно читает другой поток или нужно прервать вычисление и вернуть исходные данные, это плохой выбор. Версия с $O(m + n)$ во время первого прохода матрицу не трогает вообще.

Где ещё работает этот приём

Идея «использовать сам вход как вспомогательную память» встречается регулярно. Она хороша, когда часть данных после обработки всё равно будет перезаписана или когда в значениях есть свободные биты.

• Game of Life (LeetCode 289). Нужно обновить клетки на месте, а новое состояние зависит от старых соседей. Решение: хранить новое состояние во втором бите числа и сдвигать в конце. Это родственник нашего приёма; приёмы с битами подробно разобраны в статье о битовых операциях.
• First Missing Positive (LeetCode 41). Массив сам становится хеш-таблицей: знак числа на позиции i говорит, встречалось ли значение i + 1.
• Find All Duplicates in an Array (LeetCode 442). Та же идея со знаком.
• Product of Array Except Self. Выходной массив используют как хранилище префиксных произведений, чтобы не заводить второй. Подробно — в отдельном разборе.

Общая схема у всех одна: найти место, которое либо станет известным значением, либо может хранить лишнюю информацию без потерь, и аккуратно разобраться с границей, где служебные данные пересекаются с настоящими. В Set Matrix Zeroes такая граница — ячейка [0][0].

Типичные ошибки

• Обнулять на лету. Самая частая. Записанные нули неотличимы от исходных.
• Копировать матрицу через matrix[:] или matrix.copy(). Это поверхностная копия, строки общие. Нужно [row[:] for row in matrix].
• Один маркер в [0][0] на строку и столбец сразу. Разобрано выше: ноль в первом столбце ложно обнуляет первую строку, и наоборот.
• Обнулять первую строку и столбец до внутренней части. Маркеры превращаются в «настоящие» нули и обнуляют лишнее.
• Считать флаги после разметки. first_row_zero надо вычислить до шага 1. После него в первой строке есть нули-маркеры, и флаг почти всегда окажется истинным.
• Начинать шаги 1 и 2 с индекса 0. Тогда служебная строка и столбец смешиваются с данными.
• Забыть про матрицу из одной строки или одного столбца. В версии с двумя флагами внутренние циклы просто не выполняются, и всё решают флаги. Это стоит проверить вслух.

На собеседовании

Задачу дают на секциях по массивам и матрицам, часто вместе со спиральным обходом и поворотом матрицы. Общий обзор таких задач — в статье о задачах на массивы. Оценивают здесь не знание трюка, а умение заметить, что запись портит данные, и аккуратно разложить решение на шаги.

Хороший порядок ответа:

1. Уточните условие: на месте или можно вернуть новую матрицу, какие значения бывают, может ли матрица быть пустой.
2. Скажите, почему обнуление на лету не работает, и покажите это на маленьком примере. Это сразу показывает, что вы видите суть задачи.
3. Предложите решение с двумя массивами флагов: $O(mn)$ по времени, $O(m + n)$ по памяти. Напишите его.
4. Если спросят про $O(1)$, объясните, что первая строка и первый столбец могут играть роль этих массивов, и сразу назовите проблему с [0][0] и порядком шагов.
5. Прогоните пример с нулём в первом столбце и крайние случаи: одна строка, один столбец, нет нулей, ноль в углу.

Вопросы, которые могут последовать:

• «Почему нельзя меньше чем $O(mn)$ по времени?» — каждую ячейку нужно прочитать.
• «А если матрица огромная и не влезает в память?» — первый проход по строкам собирает множества строк и столбцов, второй проход переписывает данные. $O(m + n)$ флагов почти всегда помещаются, даже когда сама матрица нет.
• «А если значения — биты в упакованной матрице?» — флаги строк и столбцов тоже удобно хранить битовыми масками.

Самое важное, что стоит вынести из задачи, — правило «сначала прочитать, потом писать». Оно выручит в Game of Life, в фильтрации на месте из Move Zeroes и в любой задаче, где ответ нужно записать туда же, откуда берутся данные.