콘텐츠로 이동

Aho-Corasick 알고리즘

Aho-Corasick 알고리즘을 사용하면 텍스트에서 여러 패턴을 빠르게 검색할 수 있습니다. 패턴 문자열의 집합은 사전이라고도 부릅니다. 사전을 구성하는 문자열들의 전체 길이를 $m$으로, 알파벳의 크기를 $k$로 나타내겠습니다. 이 알고리즘은 트라이를 기반으로 유한 상태 오토마톤을 $O(m k)$ 시간에 구성한 다음, 이를 사용하여 텍스트를 처리합니다.

이 알고리즘은 1975년에 Alfred Aho와 Margaret Corasick이 제안했습니다.

트라이의 구성


단어 "Java", "Rad", "Rand", "Rau", "Raum", "Rose"를 기반으로 한 트라이입니다.
이 이미지는 [nd](https://de.wikipedia.org/wiki/Benutzer:Nd)가 제작했으며 CC BY-SA 3.0 라이선스로 배포됩니다.

형식적으로, 트라이는 루트가 있는 트리로서, 트리의 각 간선에는 어떤 문자가 레이블로 붙어 있고 한 정점에서 나가는 간선들은 서로 다른 레이블을 가집니다.

트라이의 각 정점을 루트에서 그 정점까지의 경로 위에 있는 레이블들로 이루어진 문자열과 동일시하겠습니다.

또한 각 정점에는 $\text{output}$ 플래그가 있으며, 이 플래그는 해당 정점이 사전의 패턴에 대응할 때 설정됩니다.

따라서 문자열 집합에 대한 트라이란 각 $\text{output}$ 정점이 집합의 문자열 하나에 대응하고, 역으로 집합의 각 문자열이 $\text{output}$ 정점 하나에 대응하는 트라이입니다.

이제 주어진 문자열 집합에 대한 트라이를 문자열들의 전체 길이에 대해 선형 시간에 구성하는 방법을 설명하겠습니다.

트리의 정점을 나타내는 구조체를 다음과 같이 도입합니다: 

const int K = 26;

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

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

vector<Vertex> trie(1);

여기서는 트라이를 정점 구조체 $\text{Vertex}$의 배열로 저장합니다. 각 $\text{Vertex}$는 $\text{output}$ 플래그와 배열 $\text{next}[]$ 형태의 간선들을 포함하며, 여기서 $\text{next}[i]$는 문자 $i$를 따라가서 도달하는 정점의 인덱스이고, 그러한 간선이 없으면 $-1$입니다. 처음에 트라이는 인덱스가 $0$인 정점 하나, 즉 루트만으로 이루어져 있습니다.

이제 트라이에 문자열 $s$를 추가하는 함수를 구현하겠습니다. 구현은 간단합니다. 루트 노드에서 시작하여, $s$의 문자들에 대응하는 간선이 있는 동안 그 간선들을 따라갑니다. 어떤 문자에 대한 간선이 없으면 새 정점을 생성하고 간선으로 연결합니다. 과정이 끝나면 마지막 정점에 $\text{output}$ 플래그를 표시합니다.

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;
}

이 구현은 분명히 선형 시간에 동작하며, 모든 정점이 $k$개의 링크를 저장하므로 $O(m k)$의 메모리를 사용합니다.

각 정점에서 배열 대신 map을 사용하면 메모리 사용량을 $O(m)$으로 줄일 수 있습니다. 하지만 이렇게 하면 시간 복잡도가 $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"를 바탕으로 한 Aho-Corasick 오토마톤입니다.
파란색 화살표는 접미사 링크이고, 초록색 화살표는 종결 링크입니다.

정점 $p$의 접미사 링크는 정점 $p$에 대응하는 문자열의 가장 긴 진접미사를 가리키는 간선입니다. 유일한 예외는 트라이의 루트이며, 루트의 접미사 링크는 자기 자신을 가리킵니다. 이제 오토마톤의 전이에 대한 설명을 다음과 같이 다시 쓸 수 있습니다: 트라이의 현재 정점에서 현재 문자를 사용한 전이가 없는 동안(또는 루트에 도달할 때까지) 접미사 링크를 따라갑니다.

이렇게 하여 오토마톤을 구성하는 문제를 트라이의 모든 정점에 대한 접미사 링크를 찾는 문제로 환원했습니다. 그런데 묘하게도 이 접미사 링크들은 오토마톤에서 구성한 전이를 사용하여 만들게 됩니다.

루트 정점과 그 직계 자식들의 접미사 링크는 모두 루트 정점을 가리킵니다. 트리에서 더 깊이 있는 정점 $v$의 접미사 링크는 다음과 같이 계산할 수 있습니다: $p$가 $v$의 조상이고 $c$가 $p$에서 $v$로 가는 간선에 붙은 문자라면, $p$로 이동한 뒤 그 접미사 링크를 따라가고, 그곳에서 문자 $c$로 전이합니다.

따라서 전이를 찾는 문제는 접미사 링크를 찾는 문제로 환원되었고, 접미사 링크를 찾는 문제는 루트에 더 가까운 정점들에 대한 접미사 링크와 전이를 찾는 문제로 환원되었습니다. 이렇게 하여 선형 시간에 해결할 수 있는 재귀적 의존 관계를 얻습니다.

구현으로 넘어갑시다. 이제 각 정점 $v$마다 조상 $p$와 $p$에서 $v$로 가는 간선의 문자 $pch$를 저장한다는 점에 유의하십시오. 또한 각 정점에는 접미사 링크 $\text{link}$(아직 계산되지 않았다면 $-1$)를 저장하고, 배열 $\text{go}[k]$에는 각 기호에 대한 오토마톤의 전이(마찬가지로 아직 계산되지 않았다면 $-1$)를 저장합니다.

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];
}

접미사 링크와 전이를 메모이제이션하므로 모든 접미사 링크와 전이를 찾는 데 드는 전체 시간이 선형이라는 것은 쉽게 알 수 있습니다.

이 개념을 그림으로 보려면 Stanford 슬라이드의 103번 슬라이드를 참고하십시오.

BFS 기반 구성

go와 get_link에 대한 재귀 호출로 전이와 접미사 링크를 계산하는 대신, 루트에서 시작하여 상향식으로 계산할 수도 있습니다. (사실 사전이 단 하나의 문자열로 이루어져 있으면, 익숙한 Knuth-Morris-Pratt 알고리즘을 얻게 됩니다.)

이 방법은 실행 시간이 전체 길이 $m$이 아니라 트라이의 정점 수 $n$에만 의존하므로, 앞에서 설명한 방법보다 몇 가지 장점이 있습니다. 또한 영속 배열(persistent array) 자료 구조를 사용하면 큰 알파벳에도 적용할 수 있으며, 이 경우 구성 시간이 $O(mk)$ 대신 $O(n \log k)$가 됩니다. $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}$ 정점을 계산하여 저장하면 최적화할 수 있습니다(이를 출구 링크(exit link)라고 부르기도 합니다). 이 값은 지연 계산 방식으로 선형 시간에 구할 수 있습니다. 따라서 각 정점에서 접미사 링크 경로상의 다음 표시된 정점, 즉 다음 매칭으로 $O(1)$ 시간에 이동할 수 있습니다. 따라서 각 매칭마다 $O(1)$ 시간을 사용하므로, $O(\text{len} + \text{ans})$ 복잡도를 달성합니다.

인덱스 자체는 찾지 않고 등장 횟수만 세고 싶다면, 각 정점 $v$에 대해 접미사 링크 경로에 있는 표시된 정점의 개수를 계산할 수 있습니다. 이는 전체 $O(n)$ 시간에 계산할 수 있습니다. 따라서 모든 매칭을 $O(\text{len})$에 합산할 수 있습니다.

주어진 어떤 문자열과도 매칭되지 않는, 주어진 길이의 사전순으로 가장 작은 문자열 찾기

문자열 집합과 길이 $L$이 주어집니다. 주어진 문자열을 하나도 포함하지 않는 길이 $L$의 문자열을 찾아야 하며, 그러한 문자열 중 사전순으로 가장 작은 것을 구해야 합니다.

문자열 집합에 대한 오토마톤을 구성할 수 있습니다. $\text{output}$ 정점은 집합에 속한 문자열과 매칭이 일어나는 상태라는 점을 떠올려 봅시다. 이 문제에서는 매칭을 피해야 하므로, 그러한 상태에는 들어갈 수 없습니다. 반면 그 외의 모든 정점에는 들어갈 수 있습니다. 따라서 오토마톤에서 모든 "나쁜" 정점을 삭제하고, 남은 오토마톤 그래프에서 길이 $L$인 사전순으로 가장 작은 경로를 찾습니다. 이 문제는 예를 들어 깊이 우선 탐색으로 $O(L)$에 해결할 수 있습니다.

주어진 모든 문자열을 포함하는 가장 짧은 문자열 찾기

여기서도 같은 아이디어를 사용합니다. 각 정점마다 해당 상태에서 매칭되는 문자열들을 나타내는 마스크를 저장합니다. 그러면 문제를 다음과 같이 다시 정의할 수 있습니다. 처음에 마스크가 0인 루트 상태 $(v = \text{root},~ \text{mask} = 0)$에서 출발하여, 모든 비트가 켜진 마스크 상태 $(v,~ \text{mask} = 2^n - 1)$에 도달하고자 합니다. 여기서 $n$은 집합에 있는 문자열의 개수입니다. 한 문자를 사용하여 한 상태에서 다른 상태로 전이할 때, 마스크도 그에 맞게 갱신합니다. 너비 우선 탐색을 실행하면 모든 비트가 켜진 마스크 상태 $(v,~ \text{mask} = 2^n - 1)$까지 가는 가장 짧은 경로를 찾을 수 있습니다.

$k$개의 문자열을 포함하는 길이 $L$의 사전순으로 가장 작은 문자열 찾기

이전 문제와 마찬가지로, 각 정점에 대해 그 정점에 대응하는 매칭의 개수(즉, 접미사 링크를 따라 도달할 수 있는 표시된 정점의 개수)를 계산합니다. 문제를 다음과 같이 재구성합니다. 현재 상태는 세 수로 이루어진 순서쌍 $(v,~ \text{len},~ \text{cnt})$로 결정되며, 상태 $(\text{root},~ 0,~ 0)$에서 상태 $(v,~ L,~ k)$에 도달하고자 합니다. 여기서 $v$는 임의의 정점이 될 수 있습니다. 따라서 깊이 우선 탐색을 사용하여 그러한 경로를 찾을 수 있습니다(그리고 탐색이 간선을 자연스러운 순서대로 살펴본다면, 찾은 경로는 자동으로 사전순으로 가장 작은 경로가 됩니다).

문제

참고 문헌