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

Алгоритм Ахо–Корасік дозволяє нам швидко шукати в тексті одразу декілька взірців. Множину рядків-взірців ще називають словником. Позначимо сумарну довжину рядків, з яких він складається, через $m$, а розмір алфавіту — через $k$. Алгоритм будує скінченний автомат на основі бора (префіксного дерева) за час $O(m k)$, а потім використовує його для обробки тексту.

Алгоритм запропонували Альфред Ахо і Маргарет Корасік у 1975 році.

Коли підходить цей алгоритм?

• Шукаєте в тексті одразу багато взірців (словник) за один прохід? (якщо взірець лише один → КМП або Z-функція)
• Множина взірців відома заздалегідь і її варто передобчислити у бор із суфіксними посиланнями?
• Потрібні саме точні входження взірців, а не запити про довільні підрядки тексту? (якщо потрібні підрядкові запити до одного рядка → Суфіксний автомат)

Побудова бора

Бор, побудований за словами "Java", "Rad", "Rand", "Rau", "Raum" і "Rose".
Зображення від nd розповсюджується за ліцензією CC BY-SA 3.0.

Формально, бор — це кореневе дерево, де кожне ребро позначене якоюсь літерою, а вихідні ребра вершини мають різні позначки.

Кожну вершину бора ми ототожнюватимемо з рядком, утвореним літерами на шляху від кореня до цієї вершини.

Кожна вершина також матиме прапорець $\text{output}$, який буде встановлено, якщо вершина відповідає одному зі взірців у словнику.

Відповідно, бором для множини рядків називаємо такий бор, у якому кожна вершина з прапорцем $\text{output}$ відповідає одному рядку з множини, і навпаки, кожному рядку множини відповідає одна вершина $\text{output}$.

Тепер опишемо, як побудувати бор для заданої множини рядків за час, лінійний відносно їхньої сумарної довжини.

Введемо структуру для вершин дерева:

C++

const int K = 26;

struct Vertex {
    int next[K];
    bool output = false;

    Vertex() {
        fill(begin(next), end(next), -1);
    }
};

vector<Vertex> trie(1);

Python

K = 26


class Vertex:
    def __init__(self):
        # next[i] — індекс вершини за символом i, або -1, якщо ребра немає
        self.next = [-1] * K
        self.output = False


# Спочатку бор складається лише з кореня (індекс 0)
trie = [Vertex()]

TypeScript

const K = 26;

class Vertex {
  // next[i] — індекс вершини за символом i, або -1, якщо ребра немає
  next: number[] = new Array<number>(K).fill(-1);
  output = false;
}

// Спочатку бор складається лише з кореня (індекс 0)
const trie: Vertex[] = [new Vertex()];

Go

const K = 26

type Vertex struct {
    next   [K]int // next[i] — індекс вершини за символом i, або -1, якщо ребра немає
    output bool
}

func newVertex() Vertex {
    var v Vertex
    for i := range v.next {
        v.next[i] = -1
    }
    return v
}

// Спочатку бор складається лише з кореня (індекс 0)
var trie = []Vertex{newVertex()}

Тут ми зберігаємо бор як масив структур $\text{Vertex}$. Кожна $\text{Vertex}$ містить прапорець $\text{output}$ та ребра у вигляді масиву $\text{next}[]$, де $\text{next}[i]$ — це індекс вершини, до якої ми потрапляємо, перейшовши за символом $i$, або $-1$, якщо такого ребра немає. Спочатку бор складається лише з однієї вершини — кореня — з індексом $0$.

Тепер реалізуємо функцію, яка додаватиме рядок $s$ до бора. Реалізація проста: ми починаємо з кореневої вершини і, доки існують ребра, що відповідають символам $s$, переходимо за ними. Якщо для якогось символу ребра немає, ми створюємо нову вершину і з'єднуємо її ребром. Наприкінці процесу позначаємо останню вершину прапорцем $\text{output}$.

C++

void add_string(string const& s) {
    int v = 0;
    for (char ch : s) {
        int c = ch - 'a';
        if (trie[v].next[c] == -1) {
            trie[v].next[c] = trie.size();
            trie.emplace_back();
        }
        v = trie[v].next[c];
    }
    trie[v].output = true;
}

Python

def add_string(s: str) -> None:
    v = 0
    for ch in s:
        c = ord(ch) - ord('a')
        if trie[v].next[c] == -1:
            trie[v].next[c] = len(trie)
            trie.append(Vertex())
        v = trie[v].next[c]
    trie[v].output = True

TypeScript

function addString(s: string): void {
  let v = 0;
  for (const ch of s) {
    const c = ch.charCodeAt(0) - "a".charCodeAt(0);
    if (trie[v].next[c] === -1) {
      trie[v].next[c] = trie.length;
      trie.push(new Vertex());
    }
    v = trie[v].next[c];
  }
  trie[v].output = true;
}

Go

func addString(s string) {
    v := 0
    for i := 0; i < len(s); i++ {
        c := int(s[i] - 'a')
        if trie[v].next[c] == -1 {
            trie[v].next[c] = len(trie)
            trie = append(trie, newVertex())
        }
        v = trie[v].next[c]
    }
    trie[v].output = true
}

Очевидно, що ця реалізація працює за лінійний час, а оскільки кожна вершина зберігає $k$ посилань, вона використовуватиме $O(m k)$ пам'яті.

Споживання пам'яті можна зменшити до $O(m)$, якщо в кожній вершині використовувати map замість масиву. Однак це збільшить часову складність до $O(m \log k)$.

Побудова автомата

Припустимо, що ми побудували бор для заданої множини рядків. Тепер погляньмо на нього з іншого боку. Якщо ми подивимося на будь-яку вершину, то рядок, що їй відповідає, є префіксом одного або кількох рядків множини, тож кожну вершину бора можна трактувати як позицію в одному або кількох рядках множини.

Насправді вершини бора можна трактувати як стани скінченного детермінованого автомата. З будь-якого стану ми можемо перейти — використавши якусь вхідну літеру — до інших станів, тобто до іншої позиції в множині рядків. Наприклад, якщо у словнику є лише один рядок $abc$, і ми стоїмо у вершині $ab$, то за літерою $c$ ми можемо перейти до вершини $abc$.

Отже, ребра бора можна розуміти як переходи в автоматі за відповідною літерою. Проте в автоматі нам потрібно мати переходи для кожної комбінації стану і літери. Якщо ми спробуємо здійснити перехід за літерою, а відповідного ребра в борі немає, то ми все одно маємо перейти в якийсь стан.

Точніше, припустимо, що ми перебуваємо у стані, який відповідає рядку $t$, і хочемо перейти до іншого стану за символом $c$. Якщо існує ребро, позначене цією літерою $c$, то ми можемо просто перейти за цим ребром і отримати вершину, що відповідає $t + c$. Якщо ж такого ребра немає, то, оскільки ми хочемо зберегти інваріант, що поточний стан є найдовшим частковим збігом в обробленому рядку, ми маємо знайти найдовший рядок у борі, що є власним суфіксом рядка $t$, і спробувати здійснити перехід звідти.

Наприклад, нехай бор побудовано за рядками $ab$ і $bc$, і ми зараз перебуваємо у вершині, що відповідає $ab$, яка водночас є вершиною $\text{output}$. Щоб перейти за літерою $c$, ми змушені перейти до стану, що відповідає рядку $b$, і звідти піти за ребром із літерою $c$.

Автомат Ахо–Корасік, побудований за словами "a", "ab", "bc", "bca", "c" і "caa".
Сині стрілки — це суфіксні посилання, зелені стрілки — термінальні посилання.

Суфіксне посилання для вершини $p$ — це ребро, що вказує на найдовший власний суфікс рядка, який відповідає вершині $p$. Єдиний особливий випадок — це корінь бора, чиє суфіксне посилання вказуватиме на самого себе. Тепер ми можемо переформулювати твердження про переходи в автоматі так: доки немає переходу з поточної вершини бора за поточною літерою (або доки не дійдемо до кореня), ми йдемо за суфіксним посиланням.

Таким чином ми звели задачу побудови автомата до задачі пошуку суфіксних посилань для всіх вершин бора. Проте, як не дивно, ми будуватимемо ці суфіксні посилання, використовуючи переходи, побудовані в автоматі.

Суфіксні посилання кореневої вершини та всіх її безпосередніх дітей вказують на кореневу вершину. Для будь-якої вершини $v$ глибше в дереві ми можемо обчислити суфіксне посилання так: якщо $p$ — це предок $v$, а $c$ — літера, що позначає ребро з $p$ до $v$, переходимо до $p$, потім ідемо за його суфіксним посиланням і звідти здійснюємо перехід за літерою $c$.

Отже, задача пошуку переходів звелася до задачі пошуку суфіксних посилань, а задача пошуку суфіксних посилань звелася до задачі пошуку суфіксного посилання і переходу, за винятком вершин, ближчих до кореня. Тож ми маємо рекурсивну залежність, яку можемо розв'язати за лінійний час.

Перейдімо до реалізації. Зауважимо, що тепер для кожної вершини $v$ ми зберігатимемо предка $p$ і символ $pch$ ребра з $p$ до $v$. Також у кожній вершині ми зберігатимемо суфіксне посилання $\text{link}$ (або $-1$, якщо воно ще не обчислене), а в масиві $\text{go}[k]$ — переходи в автоматі для кожного символу (знову ж таки $-1$, якщо ще не обчислено).

C++

const int K = 26;

struct Vertex {
    int next[K];
    bool output = false;
    int p = -1;
    char pch;
    int link = -1;
    int go[K];

    Vertex(int p=-1, char ch='$') : p(p), pch(ch) {
        fill(begin(next), end(next), -1);
        fill(begin(go), end(go), -1);
    }
};

vector<Vertex> t(1);

void add_string(string const& s) {
    int v = 0;
    for (char ch : s) {
        int c = ch - 'a';
        if (t[v].next[c] == -1) {
            t[v].next[c] = t.size();
            t.emplace_back(v, ch);
        }
        v = t[v].next[c];
    }
    t[v].output = true;
}

int go(int v, char ch);

int get_link(int v) {
    if (t[v].link == -1) {
        if (v == 0 || t[v].p == 0)
            t[v].link = 0;
        else
            t[v].link = go(get_link(t[v].p), t[v].pch);
    }
    return t[v].link;
}

int go(int v, char ch) {
    int c = ch - 'a';
    if (t[v].go[c] == -1) {
        if (t[v].next[c] != -1)
            t[v].go[c] = t[v].next[c];
        else
            t[v].go[c] = v == 0 ? 0 : go(get_link(v), ch);
    }
    return t[v].go[c];
} 

Python

K = 26


class Vertex:
    def __init__(self, p: int = -1, ch: str = '$'):
        self.next = [-1] * K
        self.output = False
        self.p = p          # предок
        self.pch = ch       # символ ребра з предка до цієї вершини
        self.link = -1      # суфіксне посилання (-1 — ще не обчислене)
        self.go = [-1] * K  # переходи автомата (-1 — ще не обчислені)


t = [Vertex()]


def add_string(s: str) -> None:
    v = 0
    for ch in s:
        c = ord(ch) - ord('a')
        if t[v].next[c] == -1:
            t[v].next[c] = len(t)
            t.append(Vertex(v, ch))
        v = t[v].next[c]
    t[v].output = True


def get_link(v: int) -> int:
    if t[v].link == -1:
        if v == 0 or t[v].p == 0:
            t[v].link = 0
        else:
            t[v].link = go(get_link(t[v].p), t[v].pch)
    return t[v].link


def go(v: int, ch: str) -> int:
    c = ord(ch) - ord('a')
    if t[v].go[c] == -1:
        if t[v].next[c] != -1:
            t[v].go[c] = t[v].next[c]
        else:
            t[v].go[c] = 0 if v == 0 else go(get_link(v), ch)
    return t[v].go[c]

TypeScript

const K = 26;

class Vertex {
  next: number[] = new Array<number>(K).fill(-1);
  output = false;
  p: number;          // предок
  pch: string;        // символ ребра з предка до цієї вершини
  link = -1;          // суфіксне посилання (-1 — ще не обчислене)
  go: number[] = new Array<number>(K).fill(-1); // переходи автомата

  constructor(p = -1, ch = "$") {
    this.p = p;
    this.pch = ch;
  }
}

const t: Vertex[] = [new Vertex()];

function addString(s: string): void {
  let v = 0;
  for (const ch of s) {
    const c = ch.charCodeAt(0) - "a".charCodeAt(0);
    if (t[v].next[c] === -1) {
      t[v].next[c] = t.length;
      t.push(new Vertex(v, ch));
    }
    v = t[v].next[c];
  }
  t[v].output = true;
}

function getLink(v: number): number {
  if (t[v].link === -1) {
    if (v === 0 || t[v].p === 0) {
      t[v].link = 0;
    } else {
      t[v].link = go(getLink(t[v].p), t[v].pch);
    }
  }
  return t[v].link;
}

function go(v: number, ch: string): number {
  const c = ch.charCodeAt(0) - "a".charCodeAt(0);
  if (t[v].go[c] === -1) {
    if (t[v].next[c] !== -1) {
      t[v].go[c] = t[v].next[c];
    } else {
      t[v].go[c] = v === 0 ? 0 : go(getLink(v), ch);
    }
  }
  return t[v].go[c];
}

Go

const K = 26

type Vertex struct {
    next   [K]int
    output bool
    p      int   // предок
    pch    byte  // символ ребра з предка до цієї вершини
    link   int   // суфіксне посилання (-1 — ще не обчислене)
    goArr  [K]int // переходи автомата (-1 — ще не обчислені)
}

func newVertex(p int, ch byte) Vertex {
    v := Vertex{p: p, pch: ch, link: -1}
    for i := range v.next {
        v.next[i] = -1
        v.goArr[i] = -1
    }
    return v
}

var t = []Vertex{newVertex(-1, '$')}

func addString(s string) {
    v := 0
    for i := 0; i < len(s); i++ {
        c := int(s[i] - 'a')
        if t[v].next[c] == -1 {
            t[v].next[c] = len(t)
            t = append(t, newVertex(v, s[i]))
        }
        v = t[v].next[c]
    }
    t[v].output = true
}

func getLink(v int) int {
    if t[v].link == -1 {
        if v == 0 || t[v].p == 0 {
            t[v].link = 0
        } else {
            t[v].link = goTo(getLink(t[v].p), t[v].pch)
        }
    }
    return t[v].link
}

func goTo(v int, ch byte) int {
    c := int(ch - 'a')
    if t[v].goArr[c] == -1 {
        if t[v].next[c] != -1 {
            t[v].goArr[c] = t[v].next[c]
        } else if v == 0 {
            t[v].goArr[c] = 0
        } else {
            t[v].goArr[c] = goTo(getLink(v), ch)
        }
    }
    return t[v].goArr[c]
}

Легко бачити, що завдяки мемоїзації суфіксних посилань і переходів загальний час на пошук усіх суфіксних посилань і переходів буде лінійним.

Для ілюстрації цієї ідеї дивіться слайд номер 103 у слайдах Стенфорда.

Побудова на основі BFS

Замість того щоб обчислювати переходи й суфіксні посилання рекурсивними викликами go та get_link, їх можна обчислити знизу вгору, починаючи з кореня. (До речі, коли словник складається лише з одного рядка, ми отримуємо знайомий алгоритм Кнута–Морріса–Пратта.)

Цей підхід має певні переваги над описаним вище, бо замість сумарної довжини $m$ його час роботи залежить лише від кількості вершин $n$ у борі. Більше того, його можна адаптувати для великих алфавітів за допомогою структури даних «персистентний масив», зробивши таким чином час побудови $O(n \log k)$ замість $O(mk)$, що є суттєвим покращенням з огляду на те, що $m$ може сягати $n^2$.

Ми можемо міркувати індуктивно, спираючись на той факт, що BFS із кореня обходить вершини в порядку зростання довжини. Ми можемо припустити, що коли ми перебуваємо у вершині $v$, її суфіксне посилання $u = link[v]$ вже успішно обчислене, а для всіх вершин меншої довжини переходи з них теж повністю обчислені.

Припустимо, що зараз ми стоїмо у вершині $v$ і розглядаємо символ $c$. По суті, маємо два випадки:

1. $go[v][c] = -1$. У цьому разі ми можемо присвоїти $go[v][c] = go[u][c]$, що вже відоме за припущенням індукції;
2. $go[v][c] = w \neq -1$. У цьому разі ми можемо присвоїти $link[w] = go[u][c]$.

Таким чином ми витрачаємо $O(1)$ часу на кожну пару з вершини й символу, що дає час роботи $O(nk)$. Основні накладні витрати тут полягають у тому, що в першому випадку ми копіюємо багато переходів з $u$, тоді як переходи другого випадку утворюють бор і в сумі по всіх вершинах дають $n$. Щоб уникнути копіювання $go[u][c]$, ми можемо використати структуру даних «персистентний масив»: за її допомогою ми спочатку копіюємо $go[u]$ у $go[v]$, а потім лише оновлюємо значення для тих символів, для яких перехід відрізнятиметься. Це приводить до алгоритму зі складністю $O(n \log k)$.

Застосування

Знайти в тексті всі рядки із заданої множини

Нам дано множину рядків і текст. Ми маємо вивести всі входження всіх рядків із множини в заданий текст за $O(\text{len} + \text{ans})$, де $\text{len}$ — це довжина тексту, а $\text{ans}$ — розмір відповіді.

Ми будуємо автомат для цієї множини рядків. Тепер ми оброблятимемо текст літера за літерою за допомогою автомата, починаючи з кореня бора. Якщо в певний момент ми перебуваємо у стані $v$, а наступна літера — $c$, то ми переходимо до наступного стану за допомогою $\text{go}(v, c)$, тим самим або збільшуючи довжину поточного збіжного підрядка на $1$, або зменшуючи її, ідучи за суфіксним посиланням.

Як для стану $v$ дізнатися, чи є якісь збіги з рядками множини? По-перше, зрозуміло, що якщо ми стоїмо на вершині $\text{output}$, то рядок, що відповідає цій вершині, закінчується на цій позиції в тексті. Однак це аж ніяк не єдиний можливий випадок досягнення збігу: якщо ми можемо дійти до однієї чи кількох вершин $\text{output}$, рухаючись за суфіксними посиланнями, то для кожної знайденої вершини $\text{output}$ також буде відповідний збіг. Простий приклад, що демонструє цю ситуацію, можна створити за допомогою множини рядків $\{dabce, abc, bc\}$ і тексту $dabc$.

Отже, якщо в кожній вершині $\text{output}$ ми зберігаємо індекс відповідного їй рядка (або список індексів, якщо в множині є однакові рядки), то ми можемо за час $O(n)$ знайти індекси всіх рядків, що збігаються з поточним станом, просто йдучи за суфіксними посиланнями від поточної вершини до кореня. Це не найефективніший розв'язок, оскільки він дає загальну складність $O(n ~ \text{len})$. Однак це можна оптимізувати, обчисливши й зберігши найближчу вершину $\text{output}$, досяжну за суфіксними посиланнями (це іноді називають вихідним посиланням). Це значення ми можемо обчислити ліниво за лінійний час. Тож для кожної вершини ми можемо за час $O(1)$ перейти до наступної позначеної вершини на шляху суфіксних посилань, тобто до наступного збігу. Таким чином на кожен збіг ми витрачаємо $O(1)$ часу, і тому досягаємо складності $O(\text{len} + \text{ans})$.

Якщо ви хочете лише порахувати кількість входжень, а не знаходити самі індекси, ви можете для кожної вершини $v$ обчислити кількість позначених вершин на шляху суфіксних посилань. Це можна обчислити загалом за час $O(n)$. Таким чином ми можемо підсумувати всі збіги за $O(\text{len})$.

Пошук лексикографічно найменшого рядка заданої довжини, що не збігається з жодним із заданих рядків

Дано множину рядків і довжину $L$. Ми маємо знайти рядок довжини $L$, який не містить жодного з рядків, і вивести лексикографічно найменший із таких рядків.

Ми можемо побудувати автомат для множини рядків. Нагадаємо, що вершини $\text{output}$ — це стани, у яких ми маємо збіг з рядком із множини. Оскільки в цій задачі ми маємо уникати збігів, нам не дозволено входити в такі стани. З іншого боку, ми можемо входити в усі інші вершини. Тож ми видаляємо з автомата всі «погані» вершини, а в графі автомата, що залишився, знаходимо лексикографічно найменший шлях довжини $L$. Цю задачу можна розв'язати за $O(L)$, наприклад, за допомогою пошуку в глибину.

Пошук найкоротшого рядка, що містить усі задані рядки

Тут ми використовуємо ті самі ідеї. Для кожної вершини ми зберігаємо маску, яка позначає рядки, що збігаються в цьому стані. Тоді задачу можна переформулювати так: перебуваючи спочатку у стані $(v = \text{root},~ \text{mask} = 0)$, ми хочемо досягти стану $(v,~ \text{mask} = 2^n - 1)$, де $n$ — кількість рядків у множині. Коли ми переходимо з одного стану в інший за літерою, ми відповідно оновлюємо маску. Запустивши пошук у ширину, ми можемо знайти шлях до стану $(v,~ \text{mask} = 2^n - 1)$ найменшої довжини.

Пошук лексикографічно найменшого рядка довжини $L$, що містить $k$ рядків

Як і в попередній задачі, ми обчислюємо для кожної вершини кількість відповідних їй збігів (тобто кількість позначених вершин, досяжних за суфіксними посиланнями). Переформулюємо задачу: поточний стан визначається трійкою чисел $(v,~ \text{len},~ \text{cnt})$, і ми хочемо зі стану $(\text{root},~ 0,~ 0)$ досягти стану $(v,~ L,~ k)$, де $v$ може бути будь-якою вершиною. Тож ми можемо знайти такий шлях за допомогою пошуку в глибину (і якщо пошук розглядає ребра в їхньому природному порядку, то знайдений шлях автоматично буде лексикографічно найменшим).

Задачі

• UVA #11590 - Prefix Lookup
• UVA #11171 - SMS
• UVA #10679 - I Love Strings!!
• Codeforces - x-prime Substrings
• Codeforces - Frequency of String
• CodeChef - TWOSTRS

Джерела

• Stanford's CS166 - Aho-Corasick Automata (Condensed)

Відеоматеріали