Word Break: как разбить строку на слова словаря — от перебора до DP и trie

Коротко

Word Break — одна из тех задач, где динамическое программирование видно не сразу. В ней нет чисел, нет «минимума» или «максимума», только строка и словарь. Но под капотом та же схема, что в размене монет: позиция в строке играет роль суммы, слова — роль монет, а вопрос «можно ли набрать» заменяет «сколько монет нужно».

Условие

Дана строка s и список слов word_dict. Нужно выяснить, можно ли разрезать s на куски так, чтобы каждый кусок был словом из словаря. Слова можно использовать сколько угодно раз.

s = "leetcode",       dict = ["leet", "code"]                    → True
s = "applepenapple",  dict = ["apple", "pen"]                    → True
s = "catsandog",      dict = ["cats", "dog", "sand", "and", "cat"] → False

В третьем примере все буквы покрываются словами по отдельности (cats, and, dog), но непересекающегося разбиения нет: после cats остаётся andog, после cat — sandog, и оба хвоста не разрезаются.

На LeetCode это задача 139, а версия со списком всех разбиений — 140 (Word Break II).

Почему жадность не работает

Первая мысль на собеседовании — идти слева и откусывать слово, как только оно нашлось. Возьмём s = "goalspecial" и словарь ["go", "goal", "goals", "special"].

• Жадно по самому короткому: go, остаётся alspecial — ни одно слово так не начинается. Тупик.
• Жадно по самому длинному: goals, остаётся pecial. Тоже тупик.
• Правильный ответ: goal + special.

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

Шаг 1. Перебор

Функция can(start) отвечает: можно ли разбить суффикс, начинающийся с позиции start. Пробуем каждое слово, которое начинается в start, и рекурсивно проверяем остаток.

def word_break_brute(s, word_dict):
    words = set(word_dict)

    def can(start):
        if start == len(s):
            return True
        for end in range(start + 1, len(s) + 1):
            if s[start:end] in words and can(end):
                return True
        return False

    return can(0)

Словарь сразу переводим в множество: проверка in для списка линейная, для множества — в среднем константная плюс хеширование подстроки. Как это устроено, разобрано в статье про хеш-таблицы.

Код корректен, но на плохих входах умирает. Классический злой тест: s = "aaaa...ab" и словарь из a, aa, aaa и так далее. Каждый префикс из букв a разрезается, а последняя b делает ответ ложным, поэтому рекурсия обязана перебрать всё. Если в словаре есть все строки a, aa, …, a^k, где k — число букв a, вызовов получается ровно $2^k$:

Цифры получены прогоном кода выше со счётчиком вызовов. Каждый лишний символ удваивает работу.

Шаг 2. Мемоизация

Посмотрите на сигнатуру: can зависит только от start. Различных значений start всего $n + 1$, а вызовов — экспонента. Значит, одни и те же суффиксы проверяются снова и снова. Запоминаем ответы — это ровно тот приём, которому посвящена статья про мемоизацию:

from functools import lru_cache


def word_break_memo(s, word_dict):
    words = set(word_dict)

    @lru_cache(maxsize=None)
    def can(start):
        if start == len(s):
            return True
        for end in range(start + 1, len(s) + 1):
            if s[start:end] in words and can(end):
                return True
        return False

    return can(0)

Теперь каждое can(start) вычисляется один раз, а внутри — не больше $n$ итераций. Итого $O(n^2)$ проверок подстрок. Один нюанс: срез s[start:end] и его хеш стоят $O(end - start)$, так что честная оценка — $O(n^3)$ операций с символами в худшем случае. На собеседовании обычно говорят «$O(n^2)$ проверок» и уточняют про стоимость среза, если спросят.

Минус рекурсии в Python — глубина. Строка из нескольких тысяч символов упрётся в лимит рекурсии (по умолчанию 1000). Поэтому дальше перейдём к таблице.

Шаг 3. DP по префиксам

Перевернём направление. Пусть dp[i] — можно ли разбить префикс s[:i], то есть первые i символов.

• База: dp[0] = True — пустой префикс разбивается на ноль слов.
• Переход: dp[i] = True, если существует j < i, такое что dp[j] = True и s[j:i] есть в словаре. Иначе говоря, префикс до j уже разбит, а последний кусок s[j:i] — слово.
• Ответ: dp[n].

def word_break(s, word_dict):
    words = set(word_dict)
    n = len(s)
    dp = [False] * (n + 1)
    dp[0] = True
    for i in range(1, n + 1):
        for j in range(i):
            if dp[j] and s[j:i] in words:
                dp[i] = True
                break
    return dp[n]

break после первого найденного j экономит работу: нам не нужно знать, сколькими способами можно разбить префикс, достаточно одного. Порядок проверок тоже важен: сначала дешёвое dp[j], потом срез. Если dp[j] ложно, срез вообще не создаётся.

Если вы уже решали задачу о размене монет, структура должна показаться знакомой. Там dp[x] — можно ли набрать сумму x, а переход идёт через dp[x - coin]. Здесь вместо монеты — слово, а вместо вычитания — отрезание суффикса. Общая схема таких задач описана в статье про динамическое программирование для начинающих. Потренировать её на монетах можно в нашем тренажёре.

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

s = "catsanddog", словарь ["cat", "cats", "and", "sand", "dog"]. Индексы префиксов от 0 до 10.

Все остальные dp[i] ложны: например, для i = 5 (catsa) ни один суффикс, начинающийся в истинной позиции, не является словом. Истинными оказываются ровно позиции 0, 3, 4, 7, 10 — это легко проверить, запустив код. Заметьте, что dp[7] получилось через cat + sand, хотя есть и второй путь cats + and. Для ответа «да/нет» это неважно, но в Word Break II понадобятся оба.

Шаг 4. Ограничение по длине слова

Внутренний цикл перебирает все j от 0 до i, хотя слово длиннее самого длинного слова словаря заведомо не найдётся. Если максимальная длина слова — $L$, то смотреть стоит только j от i - L до i:

def word_break_fast(s, word_dict):
    words = set(word_dict)
    max_len = max(map(len, words), default=0)
    n = len(s)
    dp = [False] * (n + 1)
    dp[0] = True
    for i in range(1, n + 1):
        for j in range(max(0, i - max_len), i):
            if dp[j] and s[j:i] in words:
                dp[i] = True
                break
    return dp[n]

Число проверок падает с $O(n^2)$ до $O(n \cdot L)$, а с учётом стоимости срезов — до $O(n \cdot L^2)$. В реальных ограничениях LeetCode строка до 300 символов, слова до 20, так что выигрыш заметный, а на длинных текстах — решающий. default=0 нужен на случай пустого словаря: max от пустой последовательности бросает исключение.

Это хороший пример уточняющего вопроса, который любят задавать после базового решения: «А если строка длиной миллион, а слова короткие?» Ответ — ограничить окно.

Шаг 5. BFS по позициям

На ту же задачу можно посмотреть как на граф. Вершины — позиции 0…n в строке. Ребро из start в end есть, если s[start:end] — слово. Вопрос «можно ли разбить строку» превращается в «достижима ли вершина n из вершины 0». Это классический обход, подробнее в статье про DFS и BFS.

from collections import deque


def word_break_bfs(s, word_dict):
    words = set(word_dict)
    max_len = max(map(len, words), default=0)
    n = len(s)
    queue = deque([0])
    seen = {0}
    while queue:
        start = queue.popleft()
        if start == n:
            return True
        for end in range(start + 1, min(n, start + max_len) + 1):
            if end not in seen and s[start:end] in words:
                seen.add(end)
                queue.append(end)
    return False

Множество seen здесь выполняет ту же роль, что мемоизация: каждая позиция обрабатывается один раз. Асимптотика та же, что у DP с ограничением длины.

Зачем тогда BFS? Во-первых, он останавливается, как только дошёл до конца, и не трогает недостижимые позиции. DP по префиксам проходит все i, даже если dp[i] ложно. Во-вторых, BFS сразу даёт разбиение с минимальным числом слов: первый раз, когда мы попали в n, — это кратчайший путь. Если интервьюер спросит «а какое минимальное количество слов нужно», графовая формулировка отвечает без переделки.

Word Break II: все разбиения

Теперь нужно вернуть не True/False, а все варианты разбиения:

s = "catsanddog"
dict = ["cat", "cats", "and", "sand", "dog"]
→ ["cat sand dog", "cats and dog"]

Главное, что нужно сказать на собеседовании сразу: ответ может быть экспоненциальным. Для s = "aaaa…a" (20 букв) и словаря ["a", "aa"] разбиений 10 946 — это число Фибоначчи, как в задаче о ступеньках. Никакой алгоритм не выведет экспоненциальный ответ за полиномиальное время, поэтому цель — не делать работы сверх размера ответа.

Мемоизация списков

Прямолинейный путь — та же рекурсия, только split(start) возвращает список разбиений суффикса:

from functools import lru_cache


def word_break_all(s, word_dict):
    words = set(word_dict)

    @lru_cache(maxsize=None)
    def split(start):
        if start == len(s):
            return [[]]
        result = []
        for end in range(start + 1, len(s) + 1):
            word = s[start:end]
            if word in words:
                for tail in split(end):
                    result.append([word] + tail)
        return result

    return [' '.join(parts) for parts in split(0)]


print(word_break_all("pineapplepenapple", ["apple", "pen", "applepen", "pine", "pineapple"]))
# ['pine apple pen apple', 'pine applepen apple', 'pineapple pen apple']

База [[]], а не []: у пустого суффикса ровно одно разбиение — пустое. Если вернуть [], ни одно слово не приклеится, и ответ всегда будет пустым. Это самая частая ошибка в этой задаче.

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

Бэктрекинг с отсечением

Второй вариант экономнее. Сначала считаем обычный DP, но с конца: good[i] — можно ли разбить суффикс s[i:]. Потом запускаем бэктрекинг, который заходит только в позиции с good[end] = True:

def word_break_all_pruned(s, word_dict):
    words = set(word_dict)
    n = len(s)
    good = [False] * (n + 1)
    good[n] = True
    for i in range(n - 1, -1, -1):
        good[i] = any(good[j] and s[i:j] in words for j in range(i + 1, n + 1))

    result = []
    path = []

    def dfs(start):
        if start == n:
            result.append(' '.join(path))
            return
        for end in range(start + 1, n + 1):
            if good[end] and s[start:end] in words:
                path.append(s[start:end])
                dfs(end)
                path.pop()

    if good[0]:
        dfs(0)
    return result

Отсечение гарантирует, что каждая ветка рекурсии заканчивается ответом. Без него на входе "aaa…ab" бэктрекинг перебрал бы все экспоненциальные разрезы префикса, чтобы в конце упереться в b и не выдать ни одного разбиения. С good[] проверка good[0] ложна, и рекурсия не запускается вовсе.

Ускорение через trie

Во всех версиях выше узкое место — проверка s[j:i] in words. Каждая такая проверка создаёт новую строку и хеширует её целиком. Для позиции start мы проверяем s[start:start+1], s[start:start+2], s[start:start+3]… и каждый раз заново проходим уже виденные символы.

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

def build_trie(words):
    root = {}
    for word in words:
        node = root
        for ch in word:
            node = node.setdefault(ch, {})
        node['$'] = True
    return root


def word_break_trie(s, word_dict):
    root = build_trie(word_dict)
    n = len(s)
    dp = [False] * (n + 1)
    dp[0] = True
    for start in range(n):
        if not dp[start]:
            continue
        node = root
        for end in range(start, n):
            node = node.get(s[end])
            if node is None:
                break
            if '$' in node:
                dp[end + 1] = True
    return dp[n]

Здесь DP «толкающий»: из каждой достижимой позиции start мы помечаем все позиции, куда можно дойти одним словом. Ключ '$' отмечает конец слова; символ $ в строках из букв не встречается, поэтому коллизии нет.

Сложность — $O(n \cdot L)$ шагов по trie без всяких срезов и хешей, плюс $O(\text{суммарная длина слов})$ на построение дерева. Ранний break даёт ещё больше на практике: если в словаре нет слов, начинающихся на xq, проход из позиции заканчивается на втором символе, а не на L-м.

Тот же trie ускоряет и Word Break II: в dfs вместо перебора end со срезами спускаемся по дереву и рекурсивно уходим в каждую позицию, где встретился конец слова. Такую же связку trie и перебора используют в задаче Word Search II — её родственница разобрана в статье про поиск слова в сетке.

Сравнение подходов

$W$ — суммарная длина слов словаря. В строках со срезами к каждой проверке добавляется стоимость среза, до $O(L)$.

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

• Словарь остаётся списком. s[j:i] in word_dict для списка — линейный поиск, и решение становится на порядок медленнее. Первая строка — words = set(word_dict).
• Забыли dp[0] = True. Без базы ни одна позиция не станет истинной.
• Сдвиг на единицу. dp длины n + 1, а dp[i] описывает префикс s[:i], а не символ s[i]. Путаница здесь даёт неверные срезы.
• Жадный выбор слова. Самое короткое или самое длинное подходящее слово — оба варианта ломаются на goalspecial.
• [] вместо [[]] в базе Word Break II. Пустой суффикс разбивается одним способом, а не нулём.
• Пустой словарь и max. max(map(len, [])) бросает ValueError, поэтому default=0.
• Глубокая рекурсия. Мемоизированная версия на строке из 5000 символов упадёт с RecursionError. Итеративный DP этого не боится.

Что спрашивают на собеседовании

Обычно задачу дают в формулировке LeetCode 139, и от кандидата ждут:

1. Объяснить, почему не работает жадность, — одним контрпримером.
2. Написать перебор или сразу заметить, что состояние — одна позиция, и перейти к мемоизации.
3. Довести до итеративного DP и назвать сложность, включая стоимость среза.

Дальше идут follow-up вопросы: ограничение по длине слова, вывод одного разбиения (храните для каждого i найденное j и восстановите путь с конца), минимальное число слов (BFS или dp[i] = min(dp[j] + 1), как в размене монет), все разбиения (Word Break II) и большой словарь (trie). Если вы дошли до trie и объяснили, почему он избавляет от повторного хеширования, задача закрыта полностью.

Хорошая тренировка после этой статьи — решить размен монет и мысленно сопоставить каждую строку кода с Word Break: позиция вместо суммы, слово вместо монеты. Когда эта аналогия становится очевидной, похожие задачи вроде Decode Ways или разбиения строки на палиндромы читаются с первого взгляда.