C++: Trie 구현하기

Trie를 개념만 알고, 생각해보니 직접 구현해본 적이 없었다. 알고리즘 문제 풀이를 할 때, 요긴하게 사용하기 위한 목적으로 정리해보고자 한다. 다른 곳에서 검색할 수 있는 것과 같이 동적으로 할당하는 것이 아닌, 문제의 조건에 맞게 정적으로 Trie를 할당해두고, 사용하는 방식을 다루고자 한다.

 

Trie?

트라이 (컴퓨팅) - 위키백과, 우리 모두의 백과사전

 Trie란 root를 기준으로 문자열을 저장할때, 각 노드에 어떤 문자가 존재하는지 개수는 몇 개인지 효율적으로 찾을 수 있는 자료구조이다. 알고리즘 문제에서 다루는 것은 각 노드마다 최대 26개의 알파벳의 자식 노드를 가질 때, 이를 구성하여 해결하는 문제들이 주를 이룬다.

 

구조체 정의

#define MAX_TRIE    200

struct Trie {
    Trie* child[26];
    int cnt;
}Tries[MAX_TRIE], *root;

 동적으로 할당한다면, Constructor(`Trie()`)와 Destructor(`~Trie()`)를 활용하여야 한다. 하지만 알고리즘 문제를 풀 때 동적으로 할당, 해제하는 것보다는 정적으로 문제에 필요한 최대 크기만큼 정의(`최대 단어 길이 * 최대 단어 개수`)해두고 사용하는 것이 보다 효율적이다.

 

함수 정의

newTrie()

int numOfWord;

Trie* newTrie() {
    Trie* newNode = &Tries[numOfWord++];
    memset(newNode, 0, sizeof(Trie));
    return newNode;
}

 해당 함수는 말그대로, 새로운 단어가 추가될 때, 정적으로 생선 된 Trie 구조체를 새로운 단어가 사용할 수 있게 해주는 것이다.

 

add

int add(char str[]) {
    Trie* node = root;

    for (int i = 0; str[i]; i++) {
        int nextAlpha = str[i] - 'a';
        if (!node->child[nextAlpha])
            node->child[nextAlpha] = newTrie();
        node = node->child[nextAlpha];
    }
    return ++node->cnt;
}

 단어가 추가되면, root로 부터 시작하여 해당 단어의 노드들을 추가해주고, 입력된 문자열이 존재할 때까지 반복한 다음 마지막 노드에는 해당 단어의 개수를 카운트하기 위해, `++node->cnt`를 하고 반환해준다.

 

search

int find(Trie* cur, char str[]) {
    for (int i = 0; str[i]; i++) {
        int nextAlpha = str[i] - 'a';

        if (!cur->child[nextAlpha])
            return 0;

        cur = cur->child[nextAlpha];
    }
    return cur->cnt;
}

int search(char str[]) {
    return find(root, str);
}

 특정 단어가 몇개가 있는지 탐색하는 방법은 위와 같다. 추가하는 방법과 마찬가지로, 단어의 알파벳별로 노드를 확인한다. for 문이 종료되기 전 (단어의 끝까지 탐색하지 못한 경우)에 다음 알파벳이 존재하지 않는다면 없는 단어이다.

 

한 단계 나아가기

와일드 카드가 포함된 단어 검색

int find(Trie* cur, int next, char str[]) {
    int wordCnt = 0;
    for (int i = next; str[i]; i++) {
        if (str[i] == '*' || str[i] == '?') {
            // 와일드 카드면 dfs처럼 가지를 뻗어나가면서 탐색.
            for (int j = 0; j < 26; j++) {
                if (cur->child[j])
                    wordCnt += find(cur->child[j], i + 1, str);
            }
            return wordCnt;
        }
        // 와일드 카드 아니면, 해당 단어만 탐색.
        else {
            int nextAlpha = str[i] - 'a';

            if (!cur->child[nextAlpha])
                return 0;

            cur = cur->child[nextAlpha];
        }
    }

    return cur->cnt;
}

 문제에 따라 `*`, `?`가 포함된 단어를 검색하라는 경우가 있다. 예를 들면, `ba*an*` 이런 식의 단어들이다. 해당 단어들은 `*`, `?` 자리에 어떠한 알파벳이 와도 상관없다. 따라서 위의 검색에서 DFS를 통해 모든 알파벳을 검색해주는 부분만 추가하면 된다.