문자열 탐색 (KMP, 트라이, 아호코라식)

1 : 1 문자열 탐색(KMP)

문자열 s에서 문자열 p를 찾는다고 하자.

// 시작 인데스들을 리턴
vector<int> kmp(const string& s, const string& p)
{
    int n = s.size();
    vector<int> ret, fail = makeFail(p);

    int matched = 0;
    for (int i = 0; i < n; ++i) {
        while (matched > 0 && s[i] != p[matched])
            matched = fail[matched - 1];

        if (s[i] == p[matched]) {
            matched++;
            if (matched == p.size()) {
                ret.push_back(i - matched + 1);
                matched = fail[matched - 1];
            }
        }
    }

    return ret;
}

vector<int> makeFail(const string& p) {
    int n = p.size(), matched = 0;
    vector<int> ret(n, 0);
    
    for (int i = 1; i < n; ++i) {
        while (matched > 0 && p[i] != p[matched])
            matched = ret[matched - 1];

        if (p[i] == p[matched]) ret[i] = ++matched;
        else ret[i] = 0;
    }
    return ret;
}

N : 1 문자열 탐색 (트라이)

여러개 문자열로 이루어진 문자열 집합에서, 특정 문자열 하나가 존재하는지 찾고 싶을 때 사용하는 알고리즘이다. 정렬 + 이분탐색을 사용한다면 이분탐색으로만O(N^2logN)의 시간복잡도가 걸리게된다. 하지만 트라이는 O(N).

동적할당을 통한 트라이 구현

const int LETTER = 26;
inline int childIndex(char c) {
    return c - 'a';
}

struct Node {
    Node* children[LETTER];
    bool terminal;

    Node() : terminal(false) {
        memset(children, 0, sizeof(children));
    }
    ~Node() {
        for (int i = 0; i < LETTER; ++i)
            delete children[i];
    }

    void insert(const char* key) {
        if (*key == 0) {
            terminal = true;
            return;
        }

        int chidx = childIndex(*key);
        if (children[chidx] == NULL)
            children[chidx] = new Node();
        children[chidx]->insert(key + 1);
    }

    Node* search(const char* key) {
        if (*key == 0) return this;
        int chidx = childIndex(*key);
        if (children[chidx] == NULL) return NULL;
        return children[chidx]->search(key + 1);
    }
};

아래처럼 사용하면 된다.

string s;
// 트라이 생성
Node* trie = new Node();
// 트라이에 문자열 삽입
trie->insert(s.c_str());
// 트라이에서 검색
Node* finded = trie->search(s.c_str());

동적 배열을 이용한 트라이 구현

트리는 배열로 구현할 수 있다면 배열로 구현하는게 여러모로 좋다. 대표적으로 힙 자료구조는 특정 성질 때문에 배열로 구현할 수 있고 덕분에 정말 빠르고 가볍다. 트라이도 일종의 트리인데, 삭제연산이 없기 때문에 동적 배열로 구현할 수 있다.

const int LETTER = 26;
inline int childIndex(char c) {
    return c - 'a';
}

struct Trie {
    struct Node {
        int children[LETTER];
        bool terminal;
        Node() : terminal(false) {
            memset(children, -1, sizeof(children));
        }
    };

    vector<Node> trie;
    int root;

    Trie() {
        root = newNode();
    }

    int newNode() {
        trie.push_back(Node());
        return (int)trie.size() - 1;
    }

    void insert(const string& s) {
        insert(s.c_str(), root);
    }

    void insert(const char* key, int node) {
        if (*key == 0) {
            trie[node].terminal = true;
            return;
        }
        int chidx = childIndex(*key);
        if (trie[node].children[chidx] == -1)
            trie[node].children[chidx] = newNode();

        insert(key + 1, trie[node].children[chidx]);
    }

    bool search(const string& s) {
        int idx = search(s.c_str(), root);
        if (idx == -1 || !trie[idx].terminal) return false;
        return true;
    } 

    int search(const char* key, int node) {
        if (*key == 0) return node;

        int chidx = childIndex(*key);
        if (trie[node].children[chidx] == -1) return -1;
        return search(key + 1, trie[node].children[chidx]);
    }
};

1 : N 문자열 검색 (아호 코라식)

한개의 문자열에서 동시에 여러개의 문자열을 찾는 알고리즘이다. 트라이 + KMP라고 생각하면 된다. KMP의 실패함수를 트라이를 통해 구현한다.

동적 할당을 이용한 아호코라식 구현

const int LETTER = 26;
inline int childIndex(char c) {
    return c - 'a';
}
struct Node {
    int terminal;
    Node* children[26];

    //aho corasick fail
    Node* fail;
    vector<int> res;

    Node() : terminal(-1), fail(NULL) {
        memset(children, 0, sizeof(children));
    }

    ~Node() {
        for (int i = 0; i < 26; ++i)
            delete children[i];
    }

    void insert(const char* key, int idx)
    {
        if (*key == 0) {
            this->terminal = idx;
            return;
        }

        int chidx = childIndex(*key);
        if (children[chidx] == NULL)
            children[chidx] = new Node();
        children[chidx]->insert(key + 1, idx);
    }
    // search 필요없음
};

void makeFail(Node* root)
{
    queue<Node*> q;
    root->fail = root;
    q.push(root);

    while (!q.empty())
    {
        Node* cur = q.front(); q.pop();

        for (int i = 0; i < 26; ++i)
        {
            if (cur->children[i] == NULL)
                continue;

            Node* child = cur->children[i];
            if (cur == root) {
                child->fail = root;
                if (child->terminal != -1)
                    child->res.push_back(child->terminal);
                q.push(child);
                continue;
            }

            // calculating fail Node
            Node* fail = cur->fail;
            while (fail != root && fail->children[i] == NULL)
                fail = fail->fail;
            if (fail->children[i] != NULL)
                fail = fail->children[i];
            child->fail = fail;
            // output process
            child->res = fail->res;
            if (child->terminal != -1)
                child->res.push_back(child->terminal);
            // 
            q.push(child);
        }
    }
}
// 출현인덱스, 문자열인덱스
vector<pair<int,int>> AhoCorasick(Node* root, const string& h)
{
    int n = h.length();
    vector<pair<int,int>> ret;

    Node* matched = root;
    for (int i = 0; i < n; ++i)
    {
        int temp = childIndex(h[i]);
        while (matched != root && matched->children[temp] == NULL)
            matched = matched->fail;

        if (matched->children[temp] != NULL) {
            matched = matched->children[temp];
            for (int i = 0; i < matched->res.size(); ++i)
                ret.emplace_back(i, matched->res[i]);
        }
    }

    return ret;
}

동적 배열을 이용한 아호 코라식

const int LETTER = 26;
inline int childIndex(char c) {
    return c - 'a';
}

struct Node {
    int children[LETTER];
    int terminal; // bool -> int(어떤 문자열인지 저장)
    vector<int> res;
    int fail;
    Node() : terminal(false), fail(-1) {
        memset(children, -1, sizeof(children));
    }
};

vector<Node> trie;

int newNode() {
    trie.push_back(Node());
    return (int)trie.size() - 1;
}

void insert(const char* key, int node, int idx) {
    if (*key == 0) {
    	// 어떤 문자열인지도 저장
        trie[node].terminal = idx; 
        return;
    }
    int chidx = childIndex(*key);
    if (trie[node].children[chidx] == -1)
        trie[node].children[chidx] = newNode();

    insert(key + 1, trie[node].children[chidx], idx);
}
// bfs를 활용한 fail trie 구현
void makeFail(int root)
{
    queue<int> q;
    trie[root].fail = root;
    q.push(root);

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

        for (int i = 0; i < 26; ++i)
        {
            if (trie[here].children[i] == -1)
                continue;

            int child = trie[here].children[i];
            // calculating fail Node
            if (here == root) {
                trie[child].fail = root;
            }
            else {
                
                int f = trie[here].fail;
                while (f != root && trie[f].children[i] == -1)
                    f = trie[f].fail;
                if (trie[f].children[i] != -1)
                    f = trie[f].children[i];
                trie[child].fail = f;
            }
            // output process
            int f = trie[child].fail;
            trie[child].res = trie[f].res;
            if (trie[child].terminal != -1)
                trie[child].res.push_back(trie[child].terminal);
            q.push(child);
        }
    }
}
// 출현인덱스, 문자열인덱스
vector<pair<int,int>> AhoCorasick(int root, const string& h)
{
    int n = h.length();
    vector<pair<int,int>> ret;

    int matched = root;
    for (int i = 0; i < n; ++i)
    {
        int c = childIndex(h[i]);
        while (matched != root && trie[matched].children[c] == -1)
            matched = trie[matched].fail;

        if (trie[matched].children[c] != -1) {
            matched = trie[matched].children[c];
            for (int j = 0; j < trie[matched].res.size(); ++j)
                ret.emplace_back(i, trie[matched].res[j]);
        }
    }

    return ret;
}

동적 배열로 구현하면 높은 확률로 시간과 메모리를 절약할 수 있다. 일부 문제는 동적배열로 구현해야만 풀리는 문제도 있었다. 물론 내가 동적할당을 좀 더 효율적으로 구현하지 못했을 수도 있다.

Reference

• 백준님 강의
• 종만북
• 노란책