Что делает алгоритм Ахо-Корасик?

Ахо-Корасик — это алгоритм для поиска сразу многих строк в одном тексте.

Например:

patterns:
he
she
his
hers

text:
ushers

В тексте встречаются:

she
he
hers

Зачем нужен Ахо-Корасик?

Если искать каждую строку отдельно через KMP или Рабин-Карп, можно получить слишком много работы.

Ахо-Корасик строит один автомат для всех pattern-ов. Потом он проходит по text один раз.

Главная идея

Ахо-Корасик — это бор с дополнительными ссылками.

• tr[v][c] — переход из вершины v по символу c
• link[v] — suffix link, куда откатываться при несовпадении
• cnt[v] — сколько pattern-ов заканчивается в вершине v
• out[v] — сколько pattern-ов заканчивается в v или в вершинах по suffix links

Сначала строим бор

Берем все pattern-ы и вставляем их в бор.

Например:

he
she
his
hers

Все общие префиксы будут храниться один раз.

Реализация через массивы

Для маленьких английских букв a...z удобно использовать массив:

int tr[N][26];

Основные массивы:

const int N = 1000005;

int tr[N][26];
int link[N];
int cnt[N];
int out[N];
int idx;

Как символ перевести в число?

Если строки состоят из маленьких английских букв:

int c = ch - 'a';

Тогда:

Вставка pattern в бор

Функция insert добавляет одну строку в бор.

void insert(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
    }

    cnt[v]++;
}

В конце cnt[v]++ означает: в этой вершине заканчивается один pattern.

Что такое suffix link?

Suffix link для вершины — это ссылка на самый длинный собственный суффикс текущей строки, который тоже есть в боре.

Например, если текущая вершина соответствует строке:

she

Ее суффиксы:

he
e

Если he есть в боре, то suffix link будет вести в вершину he.

Построение suffix links

Suffix links строятся через BFS от корня.

Если мы находимся в вершине v и идем по символу c в вершину u, то:

link[u] = tr[link[v]][c]

То есть suffix link ребенка можно найти через suffix link родителя.

Построение автомата

В функции build мы делаем две вещи:

• строим suffix links
• заполняем отсутствующие переходы

После этого переход tr[v][c] всегда будет существовать, и по text можно идти очень быстро.

Код build

void build() {
    queue<int> q;

    for (int c = 0; c < 26; c++) {
        int u = tr[0][c];

        if (u != 0) {
            q.push(u);
        }
    }

    while (!q.empty()) {
        int v = q.front();
        q.pop();

        out[v] = cnt[v] + out[link[v]];

        for (int c = 0; c < 26; c++) {
            int u = tr[v][c];

            if (u != 0) {
                link[u] = tr[link[v]][c];
                q.push(u);
            } else {
                tr[v][c] = tr[link[v]][c];
            }
        }
    }

    out[0] = cnt[0];
}

Что такое out[v]?

out[v] показывает, сколько pattern-ов заканчивается:

• прямо в вершине v
• или в вершинах, куда можно попасть по suffix links

Формула:

out[v] = cnt[v] + out[link[v]]

Если мы пришли в вершину v при обработке text, то out[v] — это количество pattern-ов, которые закончились в текущей позиции text.

Поиск количества всех вхождений

Теперь просто идем по text.

long long search(string text) {
    int v = 0;
    long long ans = 0;

    for (char ch : text) {
        int c = ch - 'a';

        v = tr[v][c];

        ans += out[v];
    }

    return ans;
}

Каждый символ text обрабатывается за O(1).

Полная программа: посчитать все вхождения

#include <bits/stdc++.h>
using namespace std;

const int N = 1000005;

int tr[N][26];
int link[N];
int cnt[N];
int out[N];
int idx;

void insert(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
    }

    cnt[v]++;
}

void build() {
    queue<int> q;

    for (int c = 0; c < 26; c++) {
        int u = tr[0][c];

        if (u != 0) {
            q.push(u);
        }
    }

    while (!q.empty()) {
        int v = q.front();
        q.pop();

        out[v] = cnt[v] + out[link[v]];

        for (int c = 0; c < 26; c++) {
            int u = tr[v][c];

            if (u != 0) {
                link[u] = tr[link[v]][c];
                q.push(u);
            } else {
                tr[v][c] = tr[link[v]][c];
            }
        }
    }

    out[0] = cnt[0];
}

long long search(string text) {
    int v = 0;
    long long ans = 0;

    for (char ch : text) {
        int c = ch - 'a';

        v = tr[v][c];

        ans += out[v];
    }

    return ans;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;

    for (int i = 0; i < n; i++) {
        string s;
        cin >> s;

        insert(s);
    }

    build();

    string text;
    cin >> text;

    cout << search(text);

    return 0;
}

Input / Output

Input

4
he
she
his
hers
ushers

Output

3

Потому что в ushers встречаются:

she
he
hers

Разбор примера

Text:

ushers

Pattern-ы:

he
she
his
hers

Когда алгоритм доходит до символа e в слове she, он находит сразу:

she
he

Потому что вершина she через suffix link связана с he.

Потом на конце hers находится еще:

hers

Как получить id найденных pattern-ов?

Если нужно не только количество, а какие именно pattern-ы нашлись, можно хранить список id в конечных вершинах.

vector<int> term[N];

При вставке:

void insert(string s, int id) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
    }

    term[v].push_back(id);
}

Простой вывод id

Во время поиска можно идти по suffix links и выводить все term. Это проще для понимания.

void searchIds(string text) {
    int v = 0;

    for (int i = 0; i < (int)text.size(); i++) {
        int c = text[i] - 'a';

        v = tr[v][c];

        int u = v;

        while (u != 0) {
            for (int id : term[u]) {
                cout << "pattern " << id << " ends at " << i << '\n';
            }

            u = link[u];
        }

        for (int id : term[0]) {
            cout << "pattern " << id << " ends at " << i << '\n';
        }
    }
}

Позиция начала pattern

Если pattern с id имеет длину len[id] и заканчивается в позиции i, то его начало:

start = i - len[id] + 1

Поэтому часто хранят массив длин:

vector<int> len;

Полная версия с id найденных pattern-ов

#include <bits/stdc++.h>
using namespace std;

const int N = 1000005;

int tr[N][26];
int link[N];
int idx;

vector<int> term[N];
vector<int> len;

void insert(string s, int id) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
    }

    term[v].push_back(id);
}

void build() {
    queue<int> q;

    for (int c = 0; c < 26; c++) {
        int u = tr[0][c];

        if (u != 0) {
            q.push(u);
        }
    }

    while (!q.empty()) {
        int v = q.front();
        q.pop();

        for (int c = 0; c < 26; c++) {
            int u = tr[v][c];

            if (u != 0) {
                link[u] = tr[link[v]][c];
                q.push(u);
            } else {
                tr[v][c] = tr[link[v]][c];
            }
        }
    }
}

void searchIds(string text) {
    int v = 0;

    for (int i = 0; i < (int)text.size(); i++) {
        int c = text[i] - 'a';

        v = tr[v][c];

        int u = v;

        while (u != 0) {
            for (int id : term[u]) {
                int start = i - len[id] + 1;

                cout << id << ' ' << start << '\n';
            }

            u = link[u];
        }

        for (int id : term[0]) {
            int start = i - len[id] + 1;

            cout << id << ' ' << start << '\n';
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;

    len.resize(n);

    for (int i = 0; i < n; i++) {
        string s;
        cin >> s;

        len[i] = s.size();

        insert(s, i);
    }

    build();

    string text;
    cin >> text;

    searchIds(text);

    return 0;
}

Если нужны только найденные pattern-ы

Можно создать массив found. Если pattern хотя бы раз встретился, отмечаем его.

vector<int> found(n);

found[id] = 1;

После обработки text выводим все id, где found[id] = 1.

Почему алгоритм быстрый?

После построения автомата каждый символ text обрабатывается одним переходом:

v = tr[v][c];

Поэтому проход по text работает за O(|text|).

Асимптотика

Здесь:

Бор для цифр

Если pattern-ы состоят из цифр 0...9, нужно поменять alphabet size на 10.

int tr[N][10];

int c = ch - '0';

Остальная логика такая же.

Если алфавит большой

Если символов очень много, массив tr[N][26] уже не подходит. Тогда можно использовать map или unordered_map.

vector<unordered_map<char, int>> tr;

Ахо-Корасик и KMP

Когда использовать Ахо-Корасик?

Типичные ошибки

• Забыть построить build() после всех insert
• Путать cnt[v] и out[v]
• Не заполнить отсутствующие переходы в build
• Использовать tr[v][c], когда символ не из нужного алфавита
• Взять слишком маленький размер массива N
• Не очистить массивы между тестами
• Забыть, что root имеет номер 0
• Делать поиск до построения suffix links

Как очистить автомат между тестами?

Если много тестов, можно чистить только созданные вершины.

void clearAho() {
    for (int i = 0; i <= idx; i++) {
        for (int c = 0; c < 26; c++) {
            tr[i][c] = 0;
        }

        link[i] = 0;
        cnt[i] = 0;
        out[i] = 0;
    }

    idx = 0;
}

Короткий шаблон

const int N = 1000005;

int tr[N][26];
int link[N];
int cnt[N];
int out[N];
int idx;

void insert(string s) {
    int v = 0;

    for (char ch : s) {
        int c = ch - 'a';

        if (tr[v][c] == 0) {
            tr[v][c] = ++idx;
        }

        v = tr[v][c];
    }

    cnt[v]++;
}

void build() {
    queue<int> q;

    for (int c = 0; c < 26; c++) {
        if (tr[0][c] != 0) {
            q.push(tr[0][c]);
        }
    }

    while (!q.empty()) {
        int v = q.front();
        q.pop();

        out[v] = cnt[v] + out[link[v]];

        for (int c = 0; c < 26; c++) {
            int u = tr[v][c];

            if (u != 0) {
                link[u] = tr[link[v]][c];
                q.push(u);
            } else {
                tr[v][c] = tr[link[v]][c];
            }
        }
    }
}

long long search(string text) {
    int v = 0;
    long long ans = 0;

    for (char ch : text) {
        int c = ch - 'a';

        v = tr[v][c];

        ans += out[v];
    }

    return ans;
}

Главное запомнить