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

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

автор Сергей Слотин
пререквизиты Префиксное дерево

Представим, что мы работаем журналистами в некотором авторитарном государстве, контролирующем СМИ, и в котором время от времени издаются законы, запрещающие упоминать определенные политические события или использовать определенные слова. Как эффективно реализовать подобную цензуру программно?

Более формально, пусть дан набор строк $s_1, s_2, \ldots, s_m$ алфавита размера $k$ суммарной длины $n$, называемый словарем, и длинный текст $t$. Необходимо определить, есть ли в тексте хотя бы одно слово из словаря, и если есть, то на какой позиции.

Мы уже умеем решать это, например, за $O(n \cdot |t|)$ независимо проверяя каждое слово на вхождение, или за $O(\max |s_i| \cdot |t|)$ хешируя все короткие подстроки. Но мы всё же хотим масштабируемое решение за $O(|t|)$, и сегодня мы построим автомат, который решит нашу задачу.

Алгоритм Ахо-Корасик (названный по фамилиям создателей, Альфреда Ахо и Маргарет Корасик; в народе — «карасик») за $O(nk)$ времени и памяти строит префиксное дерево для этого набора строк, а затем по этому дереву строит автомат, который может использоваться в этой и многих других строковых задачах. Алгоритм был открыт в 1975-м году, и с тех пор получил широкое распространение в системных программах для потоковой обработки текстов — например, в утилите grep.

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

Основная идея

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

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

  1. префиксом какой-то строки из словаря, и
  2. суффиксом считанного на данный момент текста.

Этой информации о состоянии нам достаточно: если это состояние помечено терминальным, то суффикс совпадает с какой-то строкой из словаря, и мы выводим его. Иначе по определению мы нашли самый длинный принимаемый бором суффикс, но никакая его подстрока не является запрещенной, а значит в этой позиции никакое запрещенное слово не содержится.

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

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

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

Определение. Суффиксная ссылка $l(v)$ ведёт в вершину $u \neq 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)$: если пойти по суффиксной ссылке, а затем по автоматному переходу по символу $c$, то получим собственный автоматных переход по этому символу.

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

Реализация

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

  • массив автоматных переходов go такого же размера, как и to;
  • суффиксную ссылку 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/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 — состояния более длинных строчек зависят только от менее длинных, так что можно пустить «волну» из корня — но так получается сложнее для понимания и реализации.