Plus One: прибавить единицу к массиву цифр и сложение длинных чисел

Коротко

Plus One выглядит как разминка, и обычно ею и служит. Но за десятью строками кода стоит школьное сложение столбиком, а из него вырастает целое семейство задач на длинную арифметику. Если разобраться в этой задаче один раз аккуратно, остальные решаются по шаблону.

Условие

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

Вход:  [1, 2, 3]   →  Выход: [1, 2, 4]
Вход:  [4, 3, 2, 1] → Выход: [4, 3, 2, 2]
Вход:  [9]         →  Выход: [1, 0]
Вход:  [9, 9]      →  Выход: [1, 0, 0]

На LeetCode это задача 66, в длину массив до 100 цифр. Попробовать решить самому можно в нашем тренажёре.

Сто цифр — ключевая деталь условия. Такое число не помещается ни в 32-битный, ни в 64-битный целый тип. Именно поэтому задачу дают в виде массива: проверяют, умеете ли вы работать с разрядами, а не с готовой арифметикой языка.

Наивное решение: собрать число и прибавить

Первое, что приходит в голову: склеить цифры в число, прибавить единицу и разобрать обратно.

def plus_one_int(digits):
    n = int(''.join(map(str, digits))) + 1
    return [int(c) for c in str(n)]

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

В JavaScript то же самое ломается сразу:

function plusOneNaive(digits) {
  const n = Number(digits.join('')) + 1
  return String(n).split('').map(Number)
}

const big = [9, 0, 0, 7, 1, 9, 9, 2, 5, 4, 7, 4, 0, 9, 9, 3] // 9007199254740993
console.log(plusOneNaive(big).join('')) // 9007199254740992

Результат оказался меньше входа. Number — это 64-битное число с плавающей точкой, и целые точно представляются только до Number.MAX_SAFE_INTEGER, то есть $2^{53} - 1 = 9\,007\,199\,254\,740\,991$. Шестнадцатизначное число уже за этой границей, а в задаче их бывает по сто цифр. С BigInt код снова станет правильным, но вопрос интервьюера «а без встроенной длинной арифметики?» никуда не денется. В Java и C++ нет даже такого обходного пути без BigInteger или своей реализации.

Интуиция: сложение столбиком

Вспомним, как прибавляют единицу на бумаге. Начинаем с последней цифры.

• Если она меньше 9, увеличиваем её — и всё. Остальные разряды не меняются.
• Если это 9, пишем 0 и «один в уме» переносим в следующий разряд. Там повторяется то же самое.

Отсюда главное наблюдение: перенос живёт ровно до первой цифры, которая не равна 9. Число 1299 превращается в 1300: две девятки обнулились, тройка выросла из двойки, единица даже не была прочитана.

Если девятки закончились вместе с массивом (999), перенос выходит за старший разряд. Все цифры уже стали нулями, и спереди нужна новая единица: 1000.

Других случаев нет. Это и есть полное решение.

Оптимальное решение

def plus_one(digits):
    for i in range(len(digits) - 1, -1, -1):
        if digits[i] < 9:
            digits[i] += 1
            return digits
        digits[i] = 0
    return [1] + digits

Цикл идёт с конца. Встретили не девятку — увеличили и сразу вернули массив. Встретили девятку — обнулили и пошли левее. Если цикл дошёл до конца без return, значит все цифры были девятками и теперь все стали нулями. Осталось приписать единицу спереди.

Обратите внимание: после цикла не нужно ничего проверять. Сам факт, что мы туда попали, уже доказывает, что массив состоял из девяток.

На JavaScript решение выглядит так же:

function plusOne(digits) {
  for (let i = digits.length - 1; i >= 0; i--) {
    if (digits[i] < 9) {
      digits[i]++
      return digits
    }
    digits[i] = 0
  }
  digits.unshift(1)
  return digits
}

console.log(plusOne([1, 2, 3])) // [ 1, 2, 4 ]
console.log(plusOne([9, 9])) // [ 1, 0, 0 ]

В Java массив не растёт, поэтому для случая из одних девяток создают новый массив длины $n + 1$. Java по умолчанию заполняет его нулями, так что достаточно записать единицу в нулевую ячейку: res[0] = 1. Копировать старые цифры не нужно, мы знаем, что там одни нули.

Версия с явным переносом

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

def plus_one_carry(digits):
    carry = 1
    for i in range(len(digits) - 1, -1, -1):
        total = digits[i] + carry
        digits[i] = total % 10
        carry = total // 10
        if carry == 0:
            break
    if carry:
        digits.insert(0, carry)
    return digits

total % 10 — цифра, которая остаётся в разряде, total // 10 — перенос дальше. Для единицы эти формулы дают то же самое, что ветка с девяткой: $9 + 1 = 10$, остаётся 0, переносим 1. Ранний break сохраняет главное свойство первого решения: если перенос погас, остальные цифры не трогаем.

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

Прогоним [1, 9, 9]:

На третьем шаге функция возвращает [2, 0, 0]. Для [9, 9] цикл обнулит обе цифры, не найдёт ничего меньше девятки, и после цикла получится [1] + [0, 0] = [1, 0, 0].

Тест-кейсы, которые стоит проговорить вслух на интервью:

• [0] → [1] — минимальный вход;
• [9] → [1, 0] — массив растёт из одного элемента;
• [1, 2, 3] → [1, 2, 4] — переноса нет вообще;
• [1, 9, 9] → [2, 0, 0] — перенос гаснет посередине;
• [9, 9, 9] → [1, 0, 0, 0] — перенос выходит за массив.

Сложность

Время. В худшем случае, когда все цифры девятки, проходим весь массив: $O(n)$. Плюс создание нового массива длины $n + 1$, это тоже $O(n)$.

Память. Если массив можно менять на месте, дополнительной памяти $O(1)$. Исключение — случай из одних девяток: там ответ длиннее входа, и новая память $O(n)$ неизбежна. Обычно говорят: «$O(1)$ дополнительной памяти, не считая ответа».

Средний случай. Здесь есть красивая деталь, которой можно отличиться на собеседовании. Цикл делает $k + 1$ шагов, где $k$ — число девяток в конце. Если последние цифры случайны, хвост хотя бы из $j$ девяток встречается с вероятностью $10^{-j}$. Матожидание числа шагов:

То есть в среднем Plus One работает за $O(1)$. То же верно и для последовательных прибавлений: если начать с нуля и прибавить единицу миллион раз, суммарно будет около 1,11 миллиона шагов. Каждая девятка, обнулённая переносом, сначала должна была «накопиться» девятью дешёвыми прибавлениями. Это классический пример амортизированного анализа, тот же, что у двоичного счётчика и динамического массива. Подробнее про такие оценки — в статье про нотацию O-большое.

Проверить это можно буквально за минуту:

digits = [0]
steps = 0
N = 10**6
for _ in range(N):
    i = len(digits) - 1
    while i >= 0 and digits[i] == 9:
        digits[i] = 0
        i -= 1
        steps += 1
    steps += 1
    if i < 0:
        digits.insert(0, 1)
    else:
        digits[i] += 1

print(digits, steps / N)  # [1, 0, 0, 0, 0, 0, 0] 1.111111

Обобщение: сложение длинных чисел

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

1. Идём с младших разрядов к старшим.
2. В каждом разряде складываем цифры, которые есть, и перенос.
3. В разряд записываем total % base, дальше переносим total // base.
4. Цикл продолжается, пока остались цифры хотя бы в одном числе или ненулевой перенос.

Четвёртый пункт — самое важное. Условие while i >= 0 or j >= 0 or carry одновременно обрабатывает числа разной длины и лишний старший разряд. Отдельная ветка «если остался перенос — допиши единицу» становится не нужна.

Add to Array-Form of Integer

Число в виде массива плюс обычное целое k (LeetCode 989). Трюк: можно не разбирать k на цифры, а сразу положить его в перенос.

def add_to_array_form(num, k):
    result = []
    i = len(num) - 1
    carry = k
    while i >= 0 or carry:
        if i >= 0:
            carry += num[i]
            i -= 1
        result.append(carry % 10)
        carry //= 10
    result.reverse()
    return result


print(add_to_array_form([1, 2, 0, 0], 34))  # [1, 2, 3, 4]
print(add_to_array_form([2, 1, 5], 806))  # [1, 0, 2, 1]

Перенос здесь бывает больше единицы, и это нормально: carry % 10 всё равно даёт текущую цифру, а остальное уходит в старшие разряды. Plus One — это ровно этот код при k = 1.

Цифры собираются в result от младшей к старшей, поэтому в конце один reverse(). Вставлять каждую цифру в начало списка через insert(0, ...) нельзя: каждая такая вставка стоит $O(n)$, и всё решение станет квадратичным.

Add Strings

Два числа заданы строками, встроенные большие числа использовать запрещено (LeetCode 415).

def add_strings(a, b):
    i, j = len(a) - 1, len(b) - 1
    carry = 0
    out = []
    while i >= 0 or j >= 0 or carry:
        total = carry
        if i >= 0:
            total += ord(a[i]) - ord('0')
            i -= 1
        if j >= 0:
            total += ord(b[j]) - ord('0')
            j -= 1
        out.append(chr(ord('0') + total % 10))
        carry = total // 10
    return ''.join(reversed(out))


print(add_strings('456', '77'))  # 533

Два указателя идут каждый по своей строке, короткое число просто «кончается» раньше. Цифру из символа получаем через ord(c) - ord('0'), это работает в любом языке. Больше приёмов работы со строками — в подборке задач на строки.

Add Binary и любые системы счисления

Меняем 10 на base — и тот же цикл складывает числа в любой системе счисления:

def add_in_base(a, b, base):
    i, j = len(a) - 1, len(b) - 1
    carry = 0
    out = []
    while i >= 0 or j >= 0 or carry:
        total = carry
        if i >= 0:
            total += a[i]
            i -= 1
        if j >= 0:
            total += b[j]
            j -= 1
        out.append(total % base)
        carry = total // base
    return out[::-1]


def add_binary(a, b):
    digits = add_in_base([int(c) for c in a], [int(c) for c in b], 2)
    return ''.join(map(str, digits))


print(add_binary('1011', '111'))  # 10010

Add Binary (LeetCode 67) — та же задача в двоичной системе. Есть и решение через побитовые операции: сумма без переноса — это a ^ b, перенос — (a & b) << 1, и так до исчезновения переноса. Оно разобрано в статье про битовые операции.

Add Two Numbers: перенос в связном списке

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

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next


def add_two_numbers(l1, l2):
    dummy = tail = ListNode()
    carry = 0
    while l1 or l2 or carry:
        total = carry
        if l1:
            total += l1.val
            l1 = l1.next
        if l2:
            total += l2.val
            l2 = l2.next
        tail.next = ListNode(total % 10)
        tail = tail.next
        carry = total // 10
    return dummy.next

Видно, что код повторяет Add Strings строчка в строчку, только указатели — это узлы, а не индексы. Фиктивная голова dummy избавляет от отдельного случая для первого узла.

Если же цифры лежат в списке в прямом порядке (Plus One Linked List, Add Two Numbers II), идти с конца нельзя. Обычно список разворачивают, складывают и разворачивают обратно, либо складывают цифры в стек. Общие приёмы со списками собраны в статье про задачи на связные списки.

Обратная задача: вычесть единицу

Хороший follow-up, который иногда задают после Plus One: а если вычитать? Логика зеркальная. Нули превращаются в девятки и «занимают» у старшего разряда, первая ненулевая цифра уменьшается на 1.

def minus_one(digits):
    # число больше нуля, ведущих нулей нет
    i = len(digits) - 1
    while digits[i] == 0:
        digits[i] = 9
        i -= 1
    digits[i] -= 1
    if digits[0] == 0 and len(digits) > 1:
        digits.pop(0)
    return digits


print(minus_one([1, 0, 0]))  # [9, 9]
print(minus_one([1]))  # [0]

Новый подводный камень здесь — ведущий ноль: 100 - 1 = 099. Его нужно убрать, но только если массив длиннее одной цифры, иначе пропадёт сам ноль. Цикл while без проверки границ безопасен, потому что число положительное: хотя бы одна ненулевая цифра в массиве есть.

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

Собирать число целиком. В Java, C++ и JavaScript это переполнение или потеря точности уже на 16–19 цифрах. В Python работает, но не показывает понимания задачи.

Забыть случай из одних девяток. Код обнуляет цифры, а потом возвращает [0, 0, 0] вместо [1, 0, 0, 0]. Это самая частая ошибка, и проверяется она тестом [9].

Не останавливаться после увеличения. Цикл без return или break продолжит идти влево. С флагом переноса это лишь лишние итерации, но в версиях без флага код начинает увеличивать все цифры подряд: [1, 2, 3] превращается в [2, 3, 4].

Идти не с того конца. В массиве Plus One старший разряд первый, а в Add Two Numbers — последний. Перед тем как писать цикл, проговорите, где младший разряд.

Вставка в начало в цикле. result.insert(0, digit) или unshift на каждом шаге — это $O(n^2)$. Собирайте цифры в конец и разворачивайте один раз.

Путать символ и цифру. '7' + '5' в строковом сложении даёт '75', а в Python int нужно вызывать явно. В C++ и Java цифра из символа — это c - '0'.

Забывать про перенос в условии цикла. while i >= 0 or j >= 0 без or carry теряет старший разряд: 5 + 5 превращается в 0.

Plus One на собеседовании

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

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

Возможные продолжения от интервьюера:

• прибавить не 1, а произвольное k — Add to Array-Form, перенос больше единицы;
• сложить два таких массива или две строки — Add Strings;
• то же в двоичной системе — Add Binary;
• числа в связном списке в прямом или обратном порядке — Add Two Numbers I и II, Plus One Linked List;
• вычесть единицу — обратная задача с займом и ведущим нулём;
• умножить две строки (Multiply Strings) — следующий шаг длинной арифметики: умножение столбиком в массив длины $n + m$ с переносами.

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

Итог

Plus One проверяет одну вещь: понимаете ли вы, как работает перенос. Идём с младшего разряда, девятки обнуляем, первую не девятку увеличиваем и выходим, а если девятки закончились вместе с массивом, приписываем единицу спереди.

Стоит записать то же самое через carry и условие while ... or carry, и получается универсальный шаблон длинного сложения. Им решаются Add to Array-Form, Add Strings, Add Binary и Add Two Numbers, а меняются только способ доступа к цифрам и основание системы счисления.