Удаление дубликатов из отсортированного массива с допуском до двух

Зачем это читать

На LeetCode есть задача Remove Duplicates from Sorted Array II — среднего уровня, но она может поставить в тупик даже опытных разработчиков. Многие пытаются решить её через подсчёт дубликатов, но это приводит к багам. Мы разберём элегантный подход с инвариантом, который обобщается на любое число k допустимых повторений. Вы не просто запомните код, а поймёте, как задавать правильный вопрос на каждом шаге.

Постановка задачи

Дан отсортированный массив, например [0, 0, 1, 1, 1, 1, 2, 3, 3]. Нужно преобразовать его так, чтобы каждое значение встречалось не более двух раз, а остаток массива можно игнорировать. Результат: [0, 0, 1, 1, 2, 3, 3, _, _]. Важно: модификация должна быть in-place, O(1) дополнительной памяти, и вернуть новую длину.

Ключевая идея: один вопрос

Первый и второй элементы всегда валидны — их можно сразу записывать. Начиная с индекса 2, каждый раз перед записью мы задаём вопрос: «Могу ли я записать текущее значение в позицию writePointer?» Ответ даёт сравнение с элементом на две позиции назад от writePointer. Если значения различаются — запись разрешена, иначе — пропускаем.

Почему именно две позиции? Потому что мы допускаем не более двух одинаковых значений подряд. Если бы допускалось k повторений, мы сравнивали бы с writePointer - k.

Пошаговый разбор на примере

Массив: [0, 0, 1, 1, 1, 1, 2, 3, 3]. Два указателя: readPointer (r) и writePointer (w). Начинаем оба с индекса 2.

• Шаг 1 (r=2, w=2): значение 1. Сравниваем с nums[w-2] = nums[0] = 0. 1 ≠ 0 — записываем 1 в позицию 2. r=3, w=3.
• Шаг 2 (r=3, w=3): значение 1. Сравниваем с nums[1] = 0. 1 ≠ 0 — записываем 1 в позицию 3. r=4, w=4.
• Шаг 3 (r=4, w=4): значение 1. Сравниваем с nums[2] = 1. 1 = 1 — запись запрещена. r=5, w остаётся 4.
• Шаг 4 (r=5, w=4): значение 1. Сравниваем с nums[2] = 1. Равны — пропускаем. r=6, w=4.
• Шаг 5 (r=6, w=4): значение 2. Сравниваем с nums[2] = 1. 2 ≠ 1 — записываем 2 в позицию 4. Массив: [0,0,1,1,2,1,2,3,3]. r=7, w=5.
• Шаг 6 (r=7, w=5): значение 3. Сравниваем с nums[3] = 1. 3 ≠ 1 — записываем 3 в позицию 5. r=8, w=6.
• Шаг 7 (r=8, w=6): значение 3. Сравниваем с nums[4] = 2. 3 ≠ 2 — записываем 3 в позицию 6. r=9 (конец), w=7.

Итоговая длина: 7. Первые 7 элементов: [0,0,1,1,2,3,3].

Код на Python

def remove_duplicates(nums):
    # writePointer указывает, куда записывать следующий валидный элемент
    writePointer = 0
    for readPointer in range(len(nums)):
        # Первые два элемента всегда можно записать
        # Или если текущее значение отличается от элемента на две позиции назад
        if readPointer < 2 or nums[readPointer] != nums[writePointer - 2]:
            nums[writePointer] = nums[readPointer]
            writePointer += 1
    return writePointer

# Пример использования
arr = [0, 0, 1, 1, 1, 1, 2, 3, 3]
new_len = remove_duplicates(arr)
print(arr[:new_len])  # [0, 0, 1, 1, 2, 3, 3]

Временная сложность: O(n). Дополнительная память: O(1).

Обобщение на k дубликатов

Если бы задача допускала не более k повторений, условие записи изменилось бы на readPointer < k or nums[readPointer] != nums[writePointer - k]. Всё остальное остаётся неизменным.

Практический вывод

Прямо сейчас откройте LeetCode или любой редактор и реализуйте этот алгоритм на Python. Затем измените условие на «не более трёх повторений» и убедитесь, что достаточно поменять одно число. Если вы застряли на похожей задаче — не пытайтесь угадать решение, а задайте себе вопрос: «Какой инвариант гарантирует корректность?»