Префиксное дерево (trie): когда set не хватает и как написать за 20 строк

Короткий ответ. Префиксное дерево, оно же trie, хранит слова по буквам. Общий префикс хранится один раз. Поиск слова или префикса стоит O(L), где L это длина слова. Его берут для автодополнения, проверки префиксов и задач вида Word Search II. Ниже устройство, код на Python и типичные ошибки.

Что такое префиксное дерево?

Корень дерева пустой. Каждое ребро подписано одной буквой. Путь от корня до узла это префикс. Узел с флагом «конец» значит, что здесь заканчивается слово.

Слова «кот», «код» и «кол» дают ветку «ко» и три листа. Общее начало хранится один раз.

Хеш-множество отвечает «есть ли слово». Дерево отвечает ещё и «есть ли слова с таким началом». В этом вся разница.

Когда задаче нужно префиксное дерево, а не set?

• «Автодополнение», «подсказки по первым буквам».
• «Есть ли слово с таким префиксом», «сколько слов начинается на».
• «Найти все слова из словаря в сетке букв».
• «Самый длинный общий префикс» для большого словаря.

Если в задаче нет слова «префикс» и нет перебора по буквам, хватит set. Дерево сложнее в коде, и без нужды его не пишут. Лишняя структура на собеседовании выглядит хуже, чем простая.

Как устроен узел дерева?

Узел это словарь детей и флаг конца слова. Ключ словаря это буква, значение это следующий узел.

Для алфавита из 26 букв вместо словаря берут массив из 26 ссылок. Так быстрее, но памяти уходит больше.

Иногда в узел кладут счётчик. Он хранит, сколько слов проходит через узел. Тогда вопрос «сколько слов начинается на ко» решается за два шага. Без счётчика пришлось бы обходить поддерево.

Как выглядит поиск на примере?

В дереве лежат слова «кот», «код» и «кола». Ищем «кол». Из корня идём по «к», потом по «о», потом по «л». Узел есть, но флага end на нём нет. Ответ: слова нет, а префикс есть.

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

Вставляем «кора». «к» и «о» уже есть, создаём «р» и «а». Ставим флаг end на последнем узле. Два новых узла вместо четырёх.

Как написать префиксное дерево на Python?

class Node:
    __slots__ = ("children", "end")

    def __init__(self):
        self.children = {}
        self.end = False


class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, word):
        node = self.root
        for ch in word:
            node = node.children.setdefault(ch, Node())
        node.end = True

    def _walk(self, prefix):
        node = self.root
        for ch in prefix:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

    def search(self, word):
        node = self._walk(word)
        return node is not None and node.end

    def starts_with(self, prefix):
        return self._walk(prefix) is not None

Три метода стоят O(L). Вставка создаёт не больше L новых узлов.

Как решить Word Search II через префиксное дерево?

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

Без дерева пришлось бы искать каждое слово отдельно. Это перебор для каждого слова с нуля.

С деревом делаем один обход в глубину из каждой клетки. Идём по сетке и одновременно спускаемся по дереву. Нет ребёнка с такой буквой, значит ветка тупиковая, и мы выходим. Дошли до узла с флагом end, записываем слово. Так все слова ищутся за один обход.

На каких ошибках ломается префиксное дерево?

• search возвращает True на префиксе. Забыли проверить флаг end.
• Флаг end ставят на каждом узле пути, а не на последнем.
• В Word Search II не помечают посещённые клетки и читают одну букву дважды.
• Строят дерево на каждый запрос заново вместо одного раза.
• Берут дерево там, где хватило бы set, и тратят время на лишний код.

Как объяснить приём на собеседовании?

Скажите: нужны ответы про префиксы, set их не даёт. Опишите узел: словарь детей плюс флаг. Назовите время O(L) на все три операции.

Потом сравните с сортированным списком и бинарным поиском. Он тоже находит префикс, но стоит O(log N · L), а вставка в него дорогая.

Узнавать приём среди соседних можно на странице префиксное дерево. Общая карта приёмов лежит в статье паттерны алгоритмов.

Вопросы и ответы

Сколько памяти занимает дерево?

В худшем случае столько же узлов, сколько букв во всех словах. Каждый узел это объект со словарём, то есть десятки байт. Для миллионов слов это заметно.

Как удалить слово из дерева?

Снять флаг end. Затем подниматься от листа и удалять узлы без детей и без флага. Чаще всего хватает первого шага.

Как найти все слова с данным префиксом?

Дойдите до узла префикса методом _walk. Потом обойдите его поддерево в глубину. Каждый узел с флагом end даёт слово. Для автодополнения обход обрывают после первых десяти найденных. Тогда ответ приходит быстро даже на большом словаре.

Чем отличается сжатое дерево?

Цепочку узлов с одним ребёнком склеивают в одно ребро с подстрокой. Так устроены маршрутизаторы URL. На собеседовании это спрашивают редко.