Алгоритм Ахо-Корасик

Пусть дан набор строк в алфавите размера $k$ суммарной длины $n$. Алгоритм Ахо-Корасик за $O(nk)$ времени и памяти строит префиксное дерево для этого набора строк, а затем по этому дереву строит автомат, который может использоваться в различных строковых задачах — например, для нахождения всех вхождений каждой строки из данного набора в произвольный текст за линейное время.

Алгоритм открыт в 1975-м году и получил широкое распространение в системных программах для потоковой обработки текстов, например, в утилите grep.

Автору этой статьи понадобилось по работе реализовать алгоритм Ахо-Корасик целых два раза, что на два раза больше, чем все остальные «олимпиадные» алгоритмы.

Префиксное дерево

Префиксное дерево или бор (англ. trie) — это структура данных для компактного хранения строк. Она устроена в виде дерева, где на рёбрах между вершинами написана символы, а некоторые вершины помечены терминальными.

Говорят, что бор принимает строку $s$, если существует такая вершина $v$, что, если выписать подряд все буквы на путях от корня до $v$, то получится строка $s$.

Бор сам по себе можно использовать для разных задач:

• Хранение строк — бор может занимать гораздо меньше места, чем массив или set строк.

• Сортировка строк — по бору можно пройтись dfs-ом и вывести все строки в лексикографическом порядке.

• Просто множество строк — как мы увидим, в него легко добавлять и удалять слова, а также проверять, есть ли слово в боре.

Реализация

Бор состоит из ссылающихся друг на друга вершин. В вершине обычно хранится следующая информация:

• терминальная ли вершина,
• ссылки на детей,
• возможно, какая-нибудь дополнительная зависящая от задачи информация о слове. Например, количество таких слов, заканчивающихся в вершине — так можно реализовать мультисет.

const int k = 26;

struct Vertex {
    Vertex* to[k] = {0};
    bool terminal = 0;
};

Vertex *root = new Vertex();

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

void add_string(string &s) {
    v = root;
    for (char c : s) {
        c -= 'a';
        if (!v->to[c]) 
            v->to[c] = new Vertex();
        v = v->to[c];
    }
    v->terminal = true;
}

Чтобы проверить, есть ли слово в боре, нужно так же пройти от корня по символам слова. Если в конце оказались в терминальной вершине — то слово есть. Если оказались в нетерминальной вершине или когда-нибудь потребовалось пройтись по несуществущей ссылке — то нет.

Удалить слово можно «лениво»: просто дойти до него и убрать флаг терминальности.

Как хранить ссылки

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

В этом случае можно придумать какой-нибудь другой способ хранить отображение из символа в ссылку. Например, хранить существующие ссылки в бинарном дереве (map), хэш-таблице (unordered_map) или просто в векторе. Они будут работать дольше, но зато потребление памяти в них будет линейным.

Суффиксные ссылки

Вернёмся к основной теме статьи.

Пусть дан набор строк $s_1,s_2,\dots ,s_n$, называемый cловарём и большой текст $t$. Необходимо найти все позиции, где строки словаря входят в текст. Для простоты дополнительно предположим, что строки из словаря не являются подстроками друг друга (позже мы увидим, что это требование избыточно).

Решать эту задачу будем следующим образом. Будем обрабатывать символы текста по одному и поддерживать наибольшую строку, являющуюся префиксом строки из словаря, и при этом также суффиксом считанного на данный момент текста. Если эта строка совпадает с каким-то $s_i$, то отметим текущий символ — в нём заканчивается какая-то строка из словаря.

Для этой задачи нам нужно как-то эффективно хранить и работать со всеми префиксами слов из словаря — для этого нам и понадобится префиксное дерево.

Добавим все слова в префиксное дерево и пометим соответствующие им вершины как терминальные. Теперь наша задача состоит в том, чтобы при добавлении очередного символа быстро находить вершину в префиксном дереве, которая соответcтвует наидленнейшему входящему в бор суффиксу нового выписанного префикса. Для этого нам понадобятся несколько вспомогательных понятий.

Анти-формализм. Чтобы не писать везде громоздкое «строка $s$, которая соответствуют вершине $v$», условимся дальше отождествлять вершину и соответствующую ей строку в боре.

Определение. Суффиксная ссылка $l(v)$ ведёт в вершину $u\ne v$, которая соответствует наидлиннейшему принемаемому бором суффиксу $v$.

Определение. Автоматный переход $\delta (v,c)$ ведёт в вершину, соответствующую минимальному принемаемому бором суффиксу строки $v+c$. Если переход и так существует в боре, то автоматный переход будет вести туда же.

Автоматные переходы — это именно то, что нам и надо в задаче: они ведут ровно в те вершины, которые соответствуют самому длинному «сматченному» суффиксу.

Заметим следующую связь суффиксных ссылок и автоматных переходов:

• $l(s_{:n})=\delta (l(s_{:n-1}),s_n)$.

• Если прямого перехода $v{\to }_{c}u$ не существует, то $\delta (v,c)=\delta (l(v),c)$.

Мы только что выразили $l$ и $\delta$ от строки длины $n$ через $l$ и $\delta$ от строк размера $(n-1)$. Значит, суффиксные ссылки и автоматные переходы можно найти динамическим программированием.

Реализация

По сравнению с префиксным деревом, нам помимо массива переходов в боре to нужно будет хранить некоторую дополнительную информацию:

• Сам массив автоматных переходов размера go.

• Суффиксную ссылку link.

• «Родительский» (последний) символ pch, который используется в формуле для суффиксной ссылки.

const int k = 26;

struct Vertex {
    Vertex *to[k] = {0}, *go[k] = {0};
    Vertex *link = 0, *p;
    int pch;
    Vertex (int _pch, Vertex *_p) { pch = _pch, p = _p; }
};

Vertex *root = new Vertex(-1, 0);

Добавление строки почти такое же:

void add_string(string &s) {
    Vertex *v = root;
    for (char _c : s) {
        int c -= int(_c - 'a');
        if (!v->to[c])
            v->to[c] = new Vertex(c, v);
        v = v->to[c];
    }
}

Подсчитывать динамики go и link будем «лениво» — введем для них две функции, которые будут мемоизировать свой результат выполнения.

// нам нужно объявить две функции, ссылающиеся друг на друга
// в C++ для этого нужно одну из них объявить только с сигнатурой перед другой
Vertex* go(Vertex *v, int c);

Vertex* link(Vertex *v) {
    if (!v->link) {
        // для строк длины меньше двух суффиксная ссылка это корень
        if (v == root || v->p == root)
            v->link = root;
        else // для остальных случаев формула работает
            v->link = go(link(v->p), v->pch);
    }
    return v->link;
}

Vertex* go(Vertex *v, int c) {
    if (!v->go[c]) {
        // если переход есть, то автоматный должен вести туда же
        if (v->to[c])
            v->go[c] = v->to[c];
        // если перехода нет из корня, то нужно сделать петлю
        else if (v == root)
            v->go[c] = root;
        else // для остальных случаев формула работает
            v->go[c] = go(link(v), c);
    }
    return v->go[c];
}

На самом деле, эффективнее и «чище» реализовывать построение автомата через bfs, но так получается немного сложнее для понимания.