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

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

Более формально, пусть дан набор строк $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.

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

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

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

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

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

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

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

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

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