Move Zeroes: нули в конец на месте и шаблон читающего и пишущего указателя

Коротко

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

Условие

Дан массив nums. Нужно переместить все нули в конец, сохранив относительный порядок остальных элементов, и сделать это на месте.

Вход:  [0, 1, 0, 3, 12]
Выход: [1, 3, 12, 0, 0]

В условии два ограничения, и решение определяют именно они. «На месте» запрещает собрать ответ в новом массиве. «Сохранить порядок» запрещает самый простой обмен с концом массива: так нули уедут назад, но 1, 3, 12 перемешаются.

На LeetCode (задача 283) функция ничего не возвращает. В нашем тренажёре массив нужно вернуть, это единственное отличие.

Наивные решения

Их полезно знать: на собеседовании можно сначала назвать простой вариант и объяснить, чем он плох.

Новый массив

def move_zeroes_copy(nums):
    result = [x for x in nums if x != 0]
    result += [0] * (len(nums) - len(result))
    nums[:] = result
    return nums

Работает за $O(n)$, но требует $O(n)$ памяти, то есть условие нарушено. Присваивание nums[:] = result тут принципиально. Вот так писать нельзя:

def move_zeroes_broken(nums):
    nums = [x for x in nums if x != 0] + [0] * nums.count(0)


data = [0, 1, 0]
move_zeroes_broken(data)
print(data)  # [0, 1, 0]

Такая строка создаёт новый список и привязывает к нему локальное имя. Список вызывающего кода остаётся прежним.

Удалять нули по одному

def move_zeroes_remove(nums):
    zeros = nums.count(0)
    for _ in range(zeros):
        nums.remove(0)
    nums.extend([0] * zeros)
    return nums

Память теперь $O(1)$, зато каждый remove ищет ноль и сдвигает весь хвост: так устроен динамический массив. При большом количестве нулей получается $O(n^2)$.

Устойчивая сортировка

def move_zeroes_sort(nums):
    nums.sort(key=lambda x: x == 0)
    return nums

Красивый трюк: False меньше True, а сортировка в Python устойчивая, поэтому порядок ненулевых сохранится. Но это $O(n \log n)$ по времени и $O(n)$ по памяти под ключи и слияния. Сортировать весь массив, чтобы отделить одно значение, избыточно.

Решение: читающий и пишущий указатель

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

• read смотрит на текущий элемент и сдвигается на каждом шаге;
• write показывает, куда положить следующий ненулевой, и сдвигается только после записи.

class Solution:
    def moveZeroes(self, nums: list[int]) -> list[int]:
        write = 0

        # Сдвигаем ненулевые элементы в начало
        for read in range(len(nums)):
            if nums[read] != 0:
                nums[write] = nums[read]
                write += 1

        # Всё, что правее write, должно стать нулями
        for i in range(write, len(nums)):
            nums[i] = 0

        return nums

То же на JavaScript, здесь хвост удобно заполнить через fill:

function moveZeroes(nums) {
  let write = 0

  for (let read = 0; read < nums.length; read++) {
    if (nums[read] !== 0) {
      nums[write] = nums[read]
      write++
    }
  }

  return nums.fill(0, write)
}

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

Возьмём [0, 1, 0, 3, 12]:

После цикла в начале лежит готовый префикс 1, 3, 12, а дальше остался мусор 3, 12. Заполняем позиции начиная с третьей нулями и получаем [1, 3, 12, 0, 0].

Почему ничего не теряется

Весь шаблон держится на двух фактах, которые верны на каждом шаге цикла:

1. nums[:write] — готовый ответ для уже прочитанной части, порядок сохранён.
2. write <= read, потому что write растёт не чаще read.

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

Сложность

Первый цикл делает $n$ шагов, второй — столько, сколько в массиве нулей. Всего не больше $2n$, то есть $O(n)$ по времени. Дополнительная память — $O(1)$, только два индекса. Если нужно освежить, как это считать, загляните в разбор O-нотации.

Вариант с обменом

Хвост можно не заполнять отдельно. Если вместо записи менять элементы местами, нули сами уезжают вправо:

def move_zeroes_swap(nums):
    write = 0
    for read in range(len(nums)):
        if nums[read] != 0:
            if read != write:
                nums[write], nums[read] = nums[read], nums[write]
            write += 1
    return nums

Один проход, и код чуть короче. Здесь появляется ещё один инвариант: всё в диапазоне nums[write:read] равно нулю. Проверка read != write отсекает обмен элемента с самим собой, который ничего не меняет, но тратит запись.

Обмен нередко называют более оптимальным, потому что в нём один цикл. Посчитаем.

Сколько записей на самом деле

У задачи на LeetCode есть уточнение: можно ли уменьшить общее число операций? Самый понятный способ ответить — посчитать записи в массив. Сделаем это честно, список-счётчиком:

class CountingList(list):
    def __init__(self, *args):
        super().__init__(*args)
        self.writes = 0

    def __setitem__(self, index, value):
        self.writes += 1
        super().__setitem__(index, value)


def overwrite(nums):
    write = 0
    for read in range(len(nums)):
        if nums[read] != 0:
            nums[write] = nums[read]
            write += 1
    for i in range(write, len(nums)):
        nums[i] = 0


def swap(nums):
    write = 0
    for read in range(len(nums)):
        if nums[read] != 0:
            if read != write:
                nums[write], nums[read] = nums[read], nums[write]
            write += 1


data = CountingList([0, 1, 0, 3, 12])
swap(data)
print(data, data.writes)  # [1, 3, 12, 0, 0] 6

Результаты для массивов из восьми элементов:

Закономерность простая:

• перезапись всегда делает ровно $n$ записей, даже если массив уже в нужном виде;
• обмен делает $2k$ записей, где $k$ — количество ненулевых элементов правее первого нуля. Только такие элементы и приходится двигать, и каждый стоит двух записей.

Отсюда правило: обмен выигрывает, когда $2k < n$, то есть когда двигать почти нечего. Если ноль стоит в начале, а дальше идут ненулевые значения, перезапись почти вдвое дешевле: $n$ против $2(n-1)$. Так что «обмен лучше» верно только для части входных данных.

Минимум записей

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

def overwrite_min(nums):
    write = 0
    for read in range(len(nums)):
        if nums[read] != 0:
            if nums[write] != nums[read]:
                nums[write] = nums[read]
            write += 1
    for i in range(write, len(nums)):
        if nums[i] != 0:
            nums[i] = 0
    return nums

Этот вариант пишет только в те позиции, где исходное значение отличается от итогового, и в каждую ровно один раз. Меньше не получится: любую такую позицию всё равно придётся изменить. Для [1, 0, 2, 0, 3, 0, 4, 0] это 5 записей вместо 8 и 6. Плата — лишнее сравнение на каждом шаге.

Когда это вообще важно

На обычном списке чисел асимптотика одна и та же. Я замерил оба варианта на миллионе элементов в CPython 3.9: разница не превышала примерно полутора раз, причём в какую сторону, зависело от данных. Выбирать вариант по скорости при таких цифрах обычно незачем.

Запись действительно становится заметной ценой в других условиях: элементы — тяжёлые структуры, которые копируются по значению; данные лежат во флеш-памяти с ограниченным ресурсом перезаписи; массив отображён в файл и каждая запись делает страницу «грязной». На собеседовании такое обсуждение показывает, что вы думаете дальше формулы $O(n)$.

Когда перезапись не подходит

В Move Zeroes заполнение хвоста работает, потому что все нули одинаковые и их можно создать заново. Но представьте похожую задачу: список задач, где выполненные нужно отправить в конец.

tasks = [
    {'id': 1, 'done': True},
    {'id': 2, 'done': False},
    {'id': 3, 'done': True},
    {'id': 4, 'done': False},
]

write = 0
for read in range(len(tasks)):
    if not tasks[read]['done']:
        tasks[write], tasks[read] = tasks[read], tasks[write]
        write += 1

print([task['id'] for task in tasks])  # [2, 4, 3, 1]

Перезапись здесь сломается: после цикла массив будет [2, 4, 3, 4]. Задача 1 потеряна, задача 4 записана дважды, и чем заполнять хвост, неизвестно. Обмен сохраняет все элементы, потому что это перестановка.

Посмотрите на результат: невыполненные идут в исходном порядке, а вот хвост стал 3, 1. Обмен устойчив только для оставляемых элементов. Если порядок нужен в обеих частях, простого решения за $O(n)$ с $O(1)$ памяти этим приёмом не получить. Проще завести второй массив либо взять готовый алгоритм вроде std::stable_partition в C++.

Шаблон: фильтрация на месте

Убираем из решения всё, что касается нулей, и остаётся общая форма:

def filter_in_place(nums, keep):
    write = 0
    for read in range(len(nums)):
        if keep(nums[read]):
            nums[write] = nums[read]
            write += 1
    return write  # новая длина: ответ лежит в nums[:write]


data = [5, 2, 8, 1, 6]
length = filter_in_place(data, lambda x: x % 2 == 0)
print(data[:length])  # [2, 8, 6]

Меняется только условие и то, что делать с хвостом. Если хвост не нужен, в Python его отрезает del nums[write:], в JavaScript — nums.length = write.

Удалить элемент

Remove Element (LeetCode 27) — та же Move Zeroes, но вместо нуля значение val, а хвост никого не интересует:

def remove_element(nums, val):
    write = 0
    for read in range(len(nums)):
        if nums[read] != val:
            nums[write] = nums[read]
            write += 1
    return write


data = [3, 2, 2, 3]
print(remove_element(data, 3), data)  # 2 [2, 2, 2, 3]

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

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

def remove_element_unordered(nums, val):
    i = 0
    n = len(nums)
    while i < n:
        if nums[i] == val:
            nums[i] = nums[n - 1]
            n -= 1  # i не двигаем: новый элемент ещё не проверен
        else:
            i += 1
    return n

Здесь записей ровно столько, сколько удалённых элементов. Для [1, 2, 3, 4, 5] и val = 1 нужна одна запись, а обычный шаблон сделает четыре. Это уже встречные указатели, тот вариант, что разобран в общей статье.

Удалить дубликаты из отсортированного массива

В Remove Duplicates (LeetCode 26) решение о том, оставлять ли элемент, зависит не только от него самого. Нужно сравнить его с последним записанным:

def remove_duplicates(nums):
    write = 0
    for read in range(len(nums)):
        if write == 0 or nums[read] != nums[write - 1]:
            nums[write] = nums[read]
            write += 1
    return write


data = [0, 0, 1, 1, 1, 2, 2, 3, 3, 4]
length = remove_duplicates(data)
print(data[:length])  # [0, 1, 2, 3, 4]

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

Оставить не больше k копий

Обобщение (при $k = 2$ это LeetCode 80) показывает, зачем держать в голове именно записанный префикс:

def remove_duplicates_k(nums, k):
    write = 0
    for read in range(len(nums)):
        if write < k or nums[read] != nums[write - k]:
            nums[write] = nums[read]
            write += 1
    return write


data = [1, 1, 1, 2, 2, 3]
length = remove_duplicates_k(data, 2)
print(data[:length])  # [1, 1, 2, 2, 3]

Логика такая: массив отсортирован, значит, если nums[write - k] равен текущему, то все последние k записанных тоже ему равны, и лимит исчерпан. Сравнивать надо именно с write - k. Вариант с nums[read - 2] выглядит похоже, но смотрит на ячейку, которую мы могли уже перезаписать. На этом же массиве он вернёт [1, 1, 2, 3] и потеряет двойку.

Сводка

Последняя строка — та же идея внутри быстрой сортировки: схема Ломуто проходит массив читающим указателем и обменом собирает слева элементы меньше опорного. Теперь понятно и то, почему она неустойчива: обмен перемешивает отброшенные элементы, как с задачами выше.

Шаблон живёт и в стандартных библиотеках. std::remove_if и std::unique в C++ сдвигают нужное в начало и возвращают итератор на новый конец, то есть по сути write. Vec::retain в Rust тоже фильтрует на месте.

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

Забыть про хвост. Без второго цикла из примера получится [1, 3, 12, 3, 12]. Ошибка легко ловится прогоном на бумаге.

Менять длину массива в цикле по нему. Такой код пропускает элементы:

nums = [0, 0, 1]
for x in nums:
    if x == 0:
        nums.remove(x)

print(nums)  # [0, 1]

После удаления элементы сдвигаются влево, а итератор идёт дальше и один из них не видит. В шаблоне длина не меняется, поэтому такой проблемы нет.

Проверять истинность вместо сравнения с нулём. if x: сработает, пока в массиве только числа. Если там встречаются False, None или пустая строка, они уедут вместе с нулями. Кроме того, в Python False == 0, так что для смешанных типов нужно решить, считать ли False нулём, и проверять тип явно. В JavaScript то же самое: if (x) отбрасывает '', null и NaN, а x !== 0 нет.

Переприсвоить nums внутри функции. Разобрали выше: снаружи ничего не изменится.

Менять write вне условия. write двигается только после записи. Если увеличивать его на каждом шаге, он совпадёт с read, и алгоритм перестанет что-либо делать.

Считать, что после фильтрации весь массив корректен. Правильный ответ лежит только в nums[:write]. Печатать и проверять надо префикс, а не весь массив.

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

Move Zeroes — лёгкая задача, поэтому её скорее дадут как разминку или как часть большего разговора о массивах (обзор соседних приёмов есть в статье о задачах на массивы). Написать рабочий код здесь мало, оценивают то, что вокруг него. Хороший порядок ответа:

1. Уточните условие: можно ли использовать дополнительную память, важен ли порядок, что с пустым массивом.
2. Назовите наивный вариант с новым массивом и его $O(n)$ по памяти.
3. Напишите шаблон с двумя указателями и проговорите инвариант write <= read.
4. Прогоните пример и крайние случаи: нет нулей, одни нули, ноль в начале и в конце.

Чего стоит ждать в продолжении:

• «Можно ли меньше записей?» — обсудите $n$ против $2k$ и вариант с проверками.
• «А если порядок неважен?» — ставим последний элемент на место удаляемого.
• «А если элементы не числа, а объекты?» — перезапись теряет данные, нужен обмен.
• «Удалите дубликаты тем же приёмом» — условие теперь смотрит на записанный префикс.

Не путайте этот шаблон с быстрым и медленным указателями в связном списке. Там скорости заданы заранее и цель другая: найти цикл или середину. Здесь write двигается только тогда, когда этого требуют данные.

Закрепить приём проще всего на практике: решите Move Zeroes в тренажёре, а затем без подсказок напишите на его основе Remove Element и Remove Duplicates. Если получится за пару минут, шаблон усвоен.