[C++] 트라이(Trie)

Coding Test/Algorithm & Data Structure

[C++] 트라이(Trie)

짱정연 2024. 10. 22. 12:11

1. 트라이(Trie)

트라이(Trie)는 어떤 문자열이 특정 문자열에 해당되는지 알아낼 때 효율적으로 사용할 수 있는 자료구조이다.

트라이는 문자열 집합 내에서 중복되는 접두사들에 대응되는 노드들이 서로 연결된 트리이다.

루트는 빈 문자열로 시작하여, 리프 노드는 문자열이 된다. 만약 위와 같은 트라이에서 'ten'이라는 단어를 찾을 때, 루트 노드에서 시작하여 t를 찾는다. 그 후 자식 노드들 중 e를 찾고, 또 n을 찾는다.

리프 노드까지 도달했다면, 해당 문자열이 우리가 찾던 문자열이 되고 탐색이 완료된다.

트라이는 모든 문자열을 O(log n)으로 탐색이 가능하여 문자열 탐색이 굉장히 빠르다는 장점이 있다.

하지만, 알파벳은 총 26개이므로 문자열을 포함하는 각 노드는 26개의 자식 노드 포인터를 가지게 되어 메모리를 많이 사용한다. 그러므로, 공간복잡도를 고려하여 문제에서 사용해야 한다.

2. C++로 트라이 구현

트리 형태를 나타내기 위해서 구조체로 각 노드를 표현해야 한다. 그리고 트라이에 문자열을 삽입하는 insert 함수와, 트라이에서 문자열을 찾는 find 함수가 필요하다.

노드 구조체
- 26개의 자식 노드 포인터와 finish 여부(리프 노드에 도달했는지)를 변수로 가짐
- 자식 노드 포인터는 O(1) 탐색을 위해 정적 배열이나 해시맵으로 관리하면 좋음
insert
- 루트부터 단어의 글자를 차례대로 탐색
- 현재 노드 자식 중 해당 자식이 있으면 이동, 없다면 새로운 자식을 생성
- 단어 삽입 후 마지막 노드에 입력된 단어 정보를 추가
find
- 루트에서 시작하여 차례대로 탐색
- 현재 노드 자식 중 해당 자식이 있다면 자식으로 이동, 없으면 탐색 실패

#include <iostream>
#include <vector>
#include <unordered_map>

using namespace std;

struct TrieNode {
    // 해시맵으로 자식 노드 포인터를 구현
    unordered_map<char, TrieNode*> children;
    // 리프 노드에 도달했는지
    bool finish;

    // 생성자
    TrieNode() : children(), finish(false) {}

    void insert(string& key, int idx) {
        // 문자열의 마지막 글자에 도달
        if (idx == key.length() - 1) finish = true;
        else {
            // 현재 탐색 중인 글자
            char next = key[idx];
            // 자식 노드에 없다면 새로운 자식 생성
            if (children[next] == nullptr) children[next] = new TrieNode;

            // 다음 글자 탐색
            children[next]->insert(key, idx+1);
        }
    }

    // 접두사로 검색되면 true 리턴
    bool findPrefix(string& key, int depth) {
        // 문자열의 마지막 글자에 도달
        if (depth == key.length()-1) return true;

        // 현재 탐색 중인 글자
        char next = key[depth];
        // 자식 노드에 없다면 탐색 실패
        if (children[next] == nullptr) return false;

        // 다음 글자 재귀 탐색
        return children[next]->findPrefix(key, depth+1);
    }

    // 완전히 일치하는 단어 단위로 찾아 true 리턴
    bool findString(string& key, int depth) {
        // 문자열의 마지막 글자에 도달 + 리프 노드에 도달 = 완전히 일치하는 단어
        if (depth == key.length()-1 && finish) return true;
        // 문자열의 마지막 글자에 도달 + 리프 노드가 아님 = 접두사
        if (depth == key.length()-1 && !finish) return false;

        char next = key[depth];
        // 현재 글자의 자식 노드가 없어 탐색 실패
        if (children[next] == nullptr) return false;

        // 다음 글자 재귀 탐색
        return children[next]->findString(key, depth+1);
    }
};

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

    vector<string> words = { "what", "when", "yours", "apple", "her", "his", "you" };

    TrieNode root;
    for (string word : words) {
        root.insert(word, 0);
    }

    string target = "wh";

    if (root.findPrefix(target, 0)) {
        cout << target << "가 검색 결과에 있습니다.\n";
    } else {
        cout << target << "가 검색 결과에 없습니다.\n";
    }

    if (root.findString(target, 0)) {
        cout << target << "가 검색 결과에 있습니다.\n";
    } else {
        cout << target << "가 검색 결과에 없습니다.\n";
    }

    return 0;
}

Reference

저작자표시