Decode Ways: сколькими способами расшифровать строку цифр

Коротко

Decode Ways выглядит как задача про строки, но строки тут ни при чём. Это Climbing Stairs с одной поправкой: шаг на одну или две ступени разрешён не всегда, а только когда соответствующий кусок строки является корректной буквой. Всё сложное в задаче сосредоточено в нулях, и именно на них ломается большинство решений, написанных с ходу.

Условие

Буквы закодированы числами: A = 1, B = 2, …, Z = 26. Сообщение превратили в строку цифр без разделителей. Нужно посчитать, сколькими способами её можно расшифровать обратно.

"12"    -> 2   ("AB" = 1 2, "L" = 12)
"226"   -> 3   ("BBF" = 2 2 6, "BZ" = 2 26, "VF" = 22 6)
"06"    -> 0   (кода 06 нет, а 0 сам по себе не буква)

На LeetCode это задача 91, уровень Medium. Ответ гарантированно помещается в 32-битное целое, длина строки до 100.

Важное уточнение из условия: "06" и "6" — разные вещи. Ведущий ноль не отбрасывается, кусок "06" не превращается в F. Об это и спотыкаются.

Сначала перебор глазами

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

def all_decodings(s):
    res = []

    def go(i, path):
        if i == len(s):
            res.append("".join(path))
            return
        if s[i] == "0":
            return
        for length in (1, 2):
            chunk = s[i : i + length]
            if len(chunk) == length and 1 <= int(chunk) <= 26:
                path.append(chr(ord("A") + int(chunk) - 1))
                go(i + length, path)
                path.pop()

    go(0, [])
    return res


print(all_decodings("226"))  # ['BBF', 'BZ', 'VF']
print(all_decodings("11106"))  # ['AAJF', 'KJF']

Это обычный бэктрекинг. Обратите внимание на "11106": кажется, что вариантов должно быть много, а их всего два. Ноль намертво приклеивается к единице перед ним (10 = J), и это отрезает почти все ветки.

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

Рекурсия: считаем, а не перечисляем

Пусть go(i) — число способов расшифровать суффикс s[i:]. Тогда:

• если i == len(s), суффикс пустой, способ ровно один: ничего не делать;
• если s[i] == '0', ни одна буква с нуля не начинается, способов ноль;
• иначе откусываем одну цифру и прибавляем go(i + 1), а если следующие две цифры образуют число от 10 до 26, прибавляем ещё go(i + 2).

def num_decodings_brute(s):
    def go(i):
        if i == len(s):
            return 1
        if s[i] == "0":
            return 0
        ways = go(i + 1)
        if i + 1 < len(s) and 10 <= int(s[i : i + 2]) <= 26:
            ways += go(i + 2)
        return ways

    return go(0)


print(num_decodings_brute("226"))  # 3
print(num_decodings_brute("06"))  # 0

Проверка s[i] == '0' стоит раньше всего остального и делает сразу две вещи. Она запрещает одиночный ноль и запрещает двузначный код с ведущим нулём. Поэтому нижнюю границу 10 в проверке пары можно было бы и не писать, но с ней условие читается без подсказок.

Код правильный, но медленный. На строке из одних единиц каждый вызов порождает два новых, и дерево растёт как числа Фибоначчи. Для строки '1' * 30 функция вызывается 3 524 577 раз, а для длины 100 из условия не дождаться.

Мемоизация

Аргумент у go один — индекс от 0 до n. Значит, различных подзадач всего $n + 1$, а экспоненту дают повторные вычисления одних и тех же. Запоминаем ответы:

from functools import lru_cache


def num_decodings_memo(s):
    @lru_cache(maxsize=None)
    def go(i):
        if i == len(s):
            return 1
        if s[i] == "0":
            return 0
        ways = go(i + 1)
        if i + 1 < len(s) and 10 <= int(s[i : i + 2]) <= 26:
            ways += go(i + 2)
        return ways

    return go(0)


print(num_decodings_memo("1" * 100))  # 573147844013817084101

Теперь $O(n)$ по времени и $O(n)$ по памяти: кэш плюс стек рекурсии. Подробно сам приём разобран в статье про мемоизацию. На длине 100 стек ещё не проблема, но в Python на строках в несколько тысяч символов рекурсия упрётся в лимит глубины, так что на собеседовании стоит сразу сказать, что дальше перейдёте к итеративной версии.

Таблица снизу вверх

Перевернём направление: пусть dp[i] — число способов расшифровать префикс длины i, то есть s[:i]. Префикс заканчивается либо однобуквенным кодом (последняя цифра s[i-1]), либо двухбуквенным (s[i-2:i]). Отсюда переход:

Квадратные скобки здесь — индикатор: 1, если условие выполнено, и 0, если нет. База: dp[0] = 1, пустой префикс расшифровывается одним способом.

def num_decodings_table(s):
    n = len(s)
    dp = [0] * (n + 1)
    dp[0] = 1
    for i in range(1, n + 1):
        if s[i - 1] != "0":
            dp[i] += dp[i - 1]
        if i >= 2 and 10 <= int(s[i - 2 : i]) <= 26:
            dp[i] += dp[i - 2]
    return dp[n]


print(num_decodings_table("11106"))  # 2

Здесь проверка 10 <= ... уже обязательна: отдельной проверки на ведущий ноль пары нет, и без нижней границы '06' прошло бы как код 6.

Почему dp[0] = 1, а не 0? Возьмём строку "12". При i = 2 пара '12' подходит, и мы прибавляем dp[0]. Этот вклад означает «вся строка — одна буква L», и он должен быть равен единице. С нулём в базе ответ получился бы 1 вместо 2.

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

На шаге i = 3 видно, как ноль «съедает» варианты: из двух способов для "21" выживает только тот, где единица уходит в пару 10. Единственная расшифровка — B J A.

O(1) памяти

dp[i] зависит только от двух предыдущих значений, поэтому весь массив не нужен, хватит двух переменных. Тот же приём, что в House Robber и в числах Фибоначчи.

def num_decodings(s):
    prev2, prev1 = 0, 1  # dp[i-2], dp[i-1]
    for i, ch in enumerate(s):
        cur = prev1 if ch != "0" else 0
        if i >= 1 and 10 <= int(s[i - 1 : i + 1]) <= 26:
            cur += prev2
        prev2, prev1 = prev1, cur
    return prev1


for t in ["12", "226", "06", "10", "30", "2101", "11106"]:
    print(t, num_decodings(t))
# 12 2 / 226 3 / 06 0 / 10 1 / 30 0 / 2101 1 / 11106 2

Начальные значения — это dp[-1] и dp[0]. Первое никогда не используется (пары при i = 0 нет), поэтому его можно взять любым, ноль просто аккуратнее.

Та же функция на JavaScript, если собеседование идёт на нём:

function numDecodings(s) {
  let prev2 = 0
  let prev1 = 1
  for (let i = 0; i < s.length; i++) {
    let cur = s[i] !== '0' ? prev1 : 0
    if (i >= 1) {
      const two = Number(s.slice(i - 1, i + 1))
      if (two >= 10 && two <= 26) cur += prev2
    }
    prev2 = prev1
    prev1 = cur
  }
  return prev1
}

console.log(['12', '226', '06', '10', '30'].map(numDecodings)) // [2, 3, 0, 1, 0]

Итог: $O(n)$ по времени, $O(1)$ по дополнительной памяти. Срез s[i-1:i+1] создаёт строку из двух символов, это константа. Если хочется совсем без аллокаций, двузначное число собирается арифметикой: int(s[i-1]) * 10 + int(s[i]).

Ловушки с нулями

Ниже решение, которое выглядит правдоподобно и проходит оба примера из условия:

def num_decodings_wrong(s):
    prev2, prev1 = 1, 1
    for i in range(1, len(s)):
        cur = prev1
        if int(s[i - 1 : i + 1]) <= 26:
            cur += prev2
        prev2, prev1 = prev1, cur
    return prev1


print(num_decodings_wrong("226"))  # 3, верно
print(num_decodings_wrong("06"))  # 2, а должно быть 0

В нём сразу три ошибки, и все про ноль. Разберём, на каких строках они вылезают.

Отсюда три правила, которые стоит проговорить вслух до того, как писать код:

1. Одиночный '0' не даёт вклада. Переход dp[i-1] выполняется только для ненулевой цифры.
2. Пара начинается не с нуля. Проверяйте 10 <= x <= 26, а не x <= 26. int('06') в Python спокойно вернёт 6, и ошибку будет не видно.
3. Ноль после цифры больше 2 обнуляет ответ. 30, 40, 90 и так далее нельзя прочитать никак. Отдельно обрабатывать этот случай не нужно: оба перехода дадут ноль, и ноль протечёт до конца. Но знать, почему так, полезно, если спросят.

Хороший приём на собеседовании — сразу прогнать решение на "0", "10", "30" и "101". Четыре коротких строки покрывают почти все ошибки.

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

Связь с Climbing Stairs

В Climbing Stairs мы поднимаемся на n ступенек шагами по 1 и 2, и число способов f(n) = f(n-1) + f(n-2). В Decode Ways ступеньки — это цифры, шаг на 1 — однобуквенный код, шаг на 2 — двухбуквенный. Разница только в том, что шаг разрешён не всегда.

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

def climb(n):
    a, b = 1, 1
    for _ in range(n - 1):
        a, b = b, a + b
    return b


print([climb(n) for n in range(1, 11)])
print([num_decodings("1" * n) for n in range(1, 11)])
# оба раза: [1, 2, 3, 5, 8, 13, 21, 34, 55, 89]

Это хороший способ объяснить решение интервьюеру: «это лестница, где некоторые ступени скользкие, на них нельзя встать, и некоторые прыжки через ступень запрещены». Дальше переход пишется почти механически. Тот же шаблон «линейная динамика, зависящая от двух предыдущих состояний» встречается в House Robber, в подсчёте разбиений строки по словарю (Word Break II в версии с подсчётом) и во множестве задач про плитки. Общая схема описана в статье про динамическое программирование для начинающих.

Сложность всех версий

Здесь $\varphi \approx 1{,}618$ — золотое сечение: худший случай рекурсии, как и у наивных чисел Фибоначчи, достигается на строке без запретов, например из одних единиц. Подробнее про такие оценки — в статье про O-нотацию.

Decode Ways II: звёздочка

Частый follow-up — LeetCode 639. В строке может встречаться '*', означающий любую цифру от 1 до 9 (но не 0). Ответ берут по модулю $10^9 + 7$.

Структура та же, меняется только вес переходов. Вместо «подходит или нет» считаем, сколькими способами можно выбрать цифры:

• одна позиция: '*' даёт 9 вариантов, '0' — 0, любая другая цифра — 1;
• пара '**': коды 11–19 и 21–26, всего 15;
• пара '*d': если d <= 6, подходят 1d и 2d (2 варианта), иначе только 1d (1 вариант);
• пара '1*' даёт 9 вариантов (11–19), '2*' — 6 (21–26), остальные 'd*' — 0.

MOD = 10**9 + 7


def num_decodings_star(s):
    def one(c):
        if c == "*":
            return 9
        return 0 if c == "0" else 1

    def two(a, b):
        if a == "*" and b == "*":
            return 15
        if a == "*":
            return 2 if b <= "6" else 1
        if b == "*":
            if a == "1":
                return 9
            if a == "2":
                return 6
            return 0
        return 1 if 10 <= int(a + b) <= 26 else 0

    prev2, prev1 = 1, one(s[0])
    for i in range(1, len(s)):
        cur = (one(s[i]) * prev1 + two(s[i - 1], s[i]) * prev2) % MOD
        prev2, prev1 = prev1, cur
    return prev1


print(num_decodings_star("*"))  # 9
print(num_decodings_star("1*"))  # 18
print(num_decodings_star("**"))  # 96

Проверить такую функцию проще всего перебором: подставить вместо каждой звёздочки все цифры 1–9, сложить ответы обычной num_decodings и сравнить на случайных коротких строках. Мы так и сделали, расхождений нет. Этот же приём — медленное, но очевидно верное решение как эталон — полезен и для основной задачи.

Как рассказывать на собеседовании

Порядок, который хорошо работает:

1. Переформулировать: «последний кусок — одна или две цифры, значит, ответ складывается из двух подзадач».
2. Сразу назвать нули и привести "06", "10", "30" как примеры. Это показывает, что вы видите подвох, ещё до кода.
3. Написать рекурсию, отметить экспоненту и повторяющиеся подзадачи.
4. Перейти к таблице, объяснить dp[0] = 1.
5. Сжать до двух переменных.
6. Прогнать вручную "2101" или "226" и граничные строки.

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

Вероятные уточняющие вопросы:

• Вернуть сами расшифровки, а не их число. Перебор из начала статьи, время пропорционально размеру ответа.
• Строка очень длинная, ответ огромный. Считать по модулю, как в Decode Ways II. В Python переполнения нет, но в Java и C++ без модуля ответ переполнит long уже на длине около 90.
• Алфавит до 99 или коды переменной длины. Переход становится суммой по всем допустимым длинам последнего куска, это уже обобщённая задача о разбиении строки, близкая к Word Break.
• Строка приходит потоком. Решение с двумя переменными уже потоковое: оно читает символы слева направо и помнит только последний.

Итого

Decode Ways — это Climbing Stairs с условиями на шаги. Шаг на одну цифру разрешён, если она не ноль. Шаг на две — если число от 10 до 26. Всё остальное — стандартный путь от рекурсии к двум переменным, который стоит отработать на этой задаче один раз, чтобы потом применять не задумываясь. А перед тем как сказать «готово», прогоните "0", "10", "30" и "101".