Ахо-Корасик: реализация

К разделу

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

1. Напоминание: какие ссылки нужны

Для каждой вершины $v$ храним:

Ссылка	Что делает	Когда нужна
`next[c]`	обычное ребро бора по символу $c$	построение, спуск
`sufLink`	вершина для наидлиннейшего собственного суффикса строки $[v]$ , являющегося префиксом какого-то слова	только при построении автоматных ссылок; после построения автомата бесполезна
`autoLink[c]`	переход автомата по символу $c$ : куда попадем, прочитав $c$ в вершине $v$	поиск
`termLink`	ближайшая терминальная вершина по цепочке суффиксных ссылок (или `nullptr`)	поиск: перечислить все слова, оканчивающиеся в текущей позиции
`terminal`	номера слов словаря, оканчивающихся в этой вершине	поиск

Правило для терминальной ссылки вершины $u$ :

если u->sufLink терминальная, то u->termLink = u->sufLink;
иначе копируем: u->termLink = u->sufLink->termLink.

Для вершин первого слоя (дети корня) терминальная ссылка всегда пустая: у строки из одного символа нет собственных суффиксов-слов.

Пример. Пусть в словаре есть слова $B$ и $A B B$ , а строка $B B$ в боре есть, но терминальной не является. Суффиксная ссылка $A B B \to B B$ ведет в нетерминальную вершину, поэтому termLink(ABB) копируется у $B B$ и ведет в $B$ . Так мы не теряем слово $B$ , хотя оно лежит «через одну» суффиксную ссылку.

2. Структура вершины

cpp

struct Node {
    std::map<char, Node*> next;      // ребра бора
    std::map<char, Node*> autoLink;  // переходы автомата
    std::vector<int> terminal;       // номера слов, оканчивающихся здесь
    Node* sufLink = nullptr;
    Node* termLink = nullptr;
    std::vector<int> is_terminal;
};

Замечания по выбору:

map, а не unordered_map. Алфавит небольшой, $\log | Σ |$ (порядка 20–30) — почти константа, unordered_map здесь только хуже.
Указатели, а не индексы. Нужен деструктор (или unique_ptr), но в нашем коде он опущен.
Терминальность — это vector<int>, а не флаг. Если в словаре есть повторяющиеся слова, одна вершина должна хранить несколько номеров. Вершина терминальна, если вектор не пуст.

3. Бор и добавление слов

cpp

class Trie {
    Node* root = new Node();
    std::set<char> abc;   // алфавит, встреченный в словаре
    int count = 0;        // сколько слов добавлено (нумерация с 0)

public:
    void insert(const std::string& s) {
        Node* v = root;
        for (char ch : s) {
            abc.insert(ch);
            if (!v->next.contains(ch)) {
                v->next[ch] = new Node();
            }
            v = v->next[ch];
        }
        v->terminal.push_back(count++);
    }
    // build() и find() ниже
};

Запоминаем алфавит в abc. Это стоит $O (\log | Σ |)$ на символ, зато мы не привязываемся ни к регистру, ни к статическим массивам: автомат работает с любыми символами (строчные, заглавные, цифры, русские буквы).
Вместо флага «терминальная» в конец слова записываем его номер. Слова нумеруются в порядке добавления.

Важная деталь: Node* v = root; — именно копия указателя. Нельзя бегать по самому root, иначе потеряем корень.

Быстро проверить бор можно отладчиком: добавить пару слов и посмотреть root->next.

4. Построение автомата `build()`

Автомат строится обходом в ширину (по слоям): сначала все вершины глубины 1, потом глубины 2 и т. д. Это гарантирует, что к моменту обработки вершины $v$ все ссылки у вершин меньшей глубины (в частности у v->sufLink) уже готовы.

4.1. Инициализация корня

Три шага:

root->sufLink = root — суффиксная ссылка корня замыкается на себя.
Суффиксные ссылки всех детей корня ведут в корень.
Строим autoLink корня по всему алфавиту abc:
- есть ребро по $c$ → autoLink[c] = next[c];
- ребра нет → autoLink[c] = root (остаемся в корне).
Если в алфавите есть буква, с которой не начинается ни одно слово (например, E), все равно нужна ссылка по ней, и она ведет в корень.

Детей корня сразу отправляем в очередь на обработку.

4.2. Обработка вершины из очереди

Берем вершину $v$ из очереди и делаем ровно то же, что для корня, но с учетом суффиксной ссылки.

Для каждого ребра $v \overset{c}{\to} u$ (ребенка $u$ ):

Суффиксная ссылка ребенка: u->sufLink = v->sufLink->autoLink[c]. То есть: встали в вершину суффиксной ссылки родителя и сделали автоматный переход по той же букве.
Терминальная ссылка ребенка (правило из п. 1).
Кладем $u$ в очередь.

Для каждой буквы $c$ из алфавита строим v->autoLink[c]:

есть ребро по $c$ → autoLink[c] = next[c];
нет ребра → autoLink[c] = v->sufLink->autoLink[c] (копируем у суффиксной ссылки).

Почему так: пусть мы в вершине, соответствующей строке $A$ , и ребра по $A$ нет. Дописав еще одну $A$ и выбросив минимальный префикс так, чтобы остаток был префиксом слова, мы снова окажемся в $A$ — это как раз переход, скопированный по суффиксной ссылке.

4.3. Код

cpp

    void build() {
        root->sufLink = root;
        std::queue<Node*> q;

        // инициализация корня
        for (auto& [ch, to] : root->next) {
            to->sufLink = root;
            q.push(to);
        }
        for (char ch : abc) {
            if (root->next.contains(ch)) {
                root->autoLink[ch] = root->next[ch];
            } else {
                root->autoLink[ch] = root;
            }
        }

        // BFS по остальным вершинам
        while (!q.empty()) {
            auto v = q.front();
            q.pop();

            // дети: суффиксные и терминальные ссылки
            for (auto& [ch, to] : v->next) {
                to->sufLink = v->sufLink->autoLink[ch];
                if (!to->sufLink->terminal.empty()) {
                    to->termLink = to->sufLink;
                } else {
                    to->termLink = to->sufLink->termLink;
                }
            }

            // автоматные переходы по всему алфавиту
            for (char ch : abc) {
                if (v->next.contains(ch)) {
                    v->autoLink[ch] = v->next[ch];
                } else {
                    v->autoLink[ch] = v->sufLink->autoLink[ch];
                }
            }
        }
    }

Замечания:

Для детей корня termLink остается nullptr — и это правильно (см. п. 1). Поэтому терминальные ссылки в инициализации корня не строим.
Порядок важен: к моменту построения autoLink для $v$ у v->sufLink (меньшей глубины) автоматные ссылки уже есть.
В цикле по детям to копируется как указатель, поэтому в очередь попадает копия адреса — все корректно.
Распаковка auto& [ch, to] : ... требует достаточно свежего стандарта C++ (17+); в старых средах вроде CodeBlocks со старым MinGW может не собраться.

5. Поиск `find()`

Идем по тексту $s$ и держим текущую вершину $v$ (сначала root). На каждом символе:

Если символа нет в алфавите abc, то он вообще не встречался в словаре: по нему не существует перехода.
Иначе делаем автоматный переход v = v->autoLink[ch].
Выписываем все слова, оканчивающиеся в этой позиции:
- номера из v->terminal (само слово, соответствующее вершине);
- затем идем по цепочке termLink: p = v->termLink; while (p) { взять p->terminal; p = p->termLink; } — это все слова, являющиеся суффиксами текущей строки.

cpp

    std::set<int> find(const std::string& s) {
        std::set<int> index;  // номера найденных слов
        Node* v = root;
        for (auto ch : s) {
            if (!abc.contains(ch)) {
                v = root;      // неизвестный символ: ни одно слово через него не проходит
                continue;
            }
            v = v->autoLink[ch];
            auto p = v;
            while (p->termLink) {
                p = p->termLink;
                auto& term = p->is_terminal;
                index.insert(term.begin(), term.end());
            }
            if (!v->is_terminal.empty()) {
                auto& term = v->is_terminal;
                index.insert(term.begin(), term.end());
            }
        }
        return index;
    }

Что здесь важно:

Неизвестный символ. В тексте может встретиться символ, которого нет в словаре: в autoLink для него ничего нет, и обращение сломает программу. Любое слово не может содержать такой символ, поэтому корректное поведение — сбросить его в корень (v = root). Для краткости написано просто continue (пропуск символа), но это неверно в общем случае: пропуск склеил бы части текста по обе стороны от «чужого» символа.
Терминальные ссылки обязательны. Если проверять только саму вершину v, часть слов теряется: например, при словах ABBA и B в тексте ABBA слово B — суффикс одной из промежуточных строк, но не лежит в $v$ .
Дубликаты номеров. Одно слово может встретиться в тексте много раз; чтобы отвечать на вопрос «входит ли слово в текст», номера складываем в std::set<int>. Если нужны позиции или количество вхождений, вместо множества пишем в вектор/счетчик.
Повторяющиеся слова словаря. Благодаря vector<int> terminal для словаря {OBA, OBA} вхождение обоих номеров находится сразу: обе копии лежат в одной терминальной вершине.

6. Сложность

Пусть $n$ — суммарная длина слов словаря, $| Σ |$ — размер алфавита, $| s |$ — длина текста.

Бор: $O (n \log | Σ |)$ .
Построение: каждая вершина обрабатывается один раз, на ней перебираем алфавит → $O (n | Σ | \log | Σ |)$ с map (с массивами вместо map было бы $O (n | Σ |)$ ).
Поиск: $O (| s | \log | Σ |)$ на переходы плюс время на перечисление ответов: прыжки по termLink стоят по одному на каждое найденное вхождение.

Эта реализация намеренно не экономит память (два map в каждой вершине): для больших данных лучше массивы индексов и unique_ptr/пул вершин.

7. Проверка на контесте: задача «Словарь»

Условие. Дан набор слов и текст, требуется определить для каждого слова, присутствует ли оно в тексте как подстрока.

Формат ввода. В первой строке дан текст (не более $10^{6}$ строчных латинских букв). Далее дано число $M$ — количество слов в словаре.

В следующих $M$ строках записаны слова (не более $30$ строчных латинских букв). Слова различны и отсортированы в лексикографическом порядке.

Суммарная длина слов в словаре не более $10^{5}$ .

Формат вывода. $M$ строк вида \t{Yes}, если слово присутствует, и \t{No} иначе

Решение через автомат:

cpp

    Trie trie;
    std::string s;
    std::cin >> s;
    size_t n = 0;
    std::cin >> n;
    for (size_t i = 0; i < n; ++i) {
        std::string w;
        std::cin >> w;
        trie.insert(w);
    }
    trie.build();
    auto res = trie.find(s);
    for (int i = 0; i < n; ++i) {
        std::cout << (res.contains(i) ? "Yes" : "No") << '\n';
    }

Именно из-за повторяющихся слов пришлось перейти от одного номера в вершине к vector<int>.

Ахо-Корасик: реализация ​

1. Напоминание: какие ссылки нужны ​

2. Структура вершины ​

3. Бор и добавление слов ​

4. Построение автомата build() ​

4.1. Инициализация корня ​

4.2. Обработка вершины из очереди ​

4.3. Код ​

5. Поиск find() ​

6. Сложность ​

7. Проверка на контесте: задача «Словарь» ​