Data Structures

1.8 Trie

개념

트라이는 문자열을 문자 단위의 경로로 저장하는 트리입니다. 루트에서 시작하여 간선에 적힌 문자를 차례로 읽으면 하나의 접두사가 됩니다. 같은 접두사를 가진 문자열은 그 접두사에 해당하는 경로를 공유합니다.

단어 목록에서 특정 단어가 있는지 확인하거나, 특정 접두사로 시작하는 단어가 있는지 찾을 때 사용합니다. 예를 들어 ca로 시작하는 단어를 찾을 때는 모든 단어의 앞부분을 비교하는 대신 루트에서 c, a 두 간선을 따라갑니다.

문자열 저장과 끝 표시

car, cat, dog를 저장하면 car와 cat은 c → a 경로를 공유하고, 마지막 문자에서 갈라집니다. 아래 그림에서 00은 루트이며, 진한 노드는 단어가 끝나는 위치입니다. 간선의 문자를 루트부터 읽습니다.

c a r t d o g 0 1 2 3 4 5 6 7
▲ [그림 1] car, cat, dog를 저장한 트라이

ca를 추가할 때는 이미 있는 c → a 경로를 따라간 뒤, 도착한 노드에 단어의 끝을 표시합니다. 새 간선을 만들 필요가 없습니다. 이 표시가 있어야 접두사 ca와 저장된 단어 ca를 구분할 수 있습니다.

c a r t d o g 0 1 2 3 4 5 6 7
▲ [그림 2] ca를 추가한 트라이

삽입과 검색

삽입은 루트에서 시작합니다. 다음 문자에 해당하는 자식이 있으면 이동하고, 없으면 새 노드와 간선을 만든 뒤 이동합니다. 마지막 문자를 처리한 노드에 끝 표시를 남깁니다. 같은 단어를 다시 넣으면 같은 경로와 끝 표시를 사용합니다.

단어 검색도 문자를 하나씩 따라갑니다. 도중에 필요한 간선이 없으면 저장되지 않은 단어입니다. 모든 문자를 따라간 뒤에는 끝 표시까지 확인합니다. 위의 첫 그림에서 car는 끝 표시가 있어 검색에 성공하지만, ca는 경로만 있으므로 단어 검색에 실패합니다.

접두사 검색은 해당 경로가 있는지만 확인합니다. 첫 그림에서도 ca 경로가 있으므로 ca로 시작하는 단어가 존재합니다. 자동 완성 후보를 모두 구하려면 해당 노드에서 아래로 탐색하여 끝 표시가 있는 경로를 모을 수 있습니다.

자식을 저장하는 방법과 복잡도

영어 소문자만 저장한다면 각 노드에 2626개 자식 칸을 둡니다. 문자 a부터 z를 00부터 2525에 대응시켜 다음 노드를 바로 찾습니다. 등장하는 문자의 종류가 많으면 필요한 자식만 맵에 저장하는 방법도 있습니다.

문자열의 길이를 LL, 저장하는 문자열 길이의 합을 SS, 만들어진 노드 수를 VV라고 합니다. 루트를 포함해 V≤S+1V\le S+1입니다. 아래는 노드마다 크기 2626의 자식 배열을 사용하는 경우입니다.

연산·저장 대상시간복잡도공간복잡도
길이 LL의 문자열 삽입O(L)O(L)새 노드 최대 O(26L)O(26L)
단어·접두사 검색O(L)O(L)추가 공간 O(1)O(1)
모든 문자열 삽입·전체 저장O(S+26V)O(S+26V)O(26V)O(26V)

2626은 고정된 상수이므로 전체 공간은 O(V)O(V)로도 나타냅니다. 다만 실제로는 노드마다 자식 칸 2626개를 저장하므로, 입력 문자열의 총길이를 확인하여 배열 크기를 정해야 합니다.

예제

문제

단어 NN개를 저장한 뒤 질의 QQ개에 답하여라. 질의 word s는 단어 s가 저장되어 있는지 묻고, 질의 prefix s는 s로 시작하는 단어가 있는지 묻는다.

입력

첫째 줄에 양의 정수 NN과 QQ가 공백을 사이에 두고 주어진다. 1≤N,Q≤1001\le N,Q\le100이다.

다음 NN개 줄에 단어가 하나씩 주어진다. 이어지는 QQ개 줄에 질의 종류와 문자열 ss가 공백을 사이에 두고 주어진다. 질의 종류는 word 또는 prefix이다.

모든 단어와 질의 문자열은 영어 소문자로 이루어져 있으며, 길이는 11 이상 2020 이하이다.

출력

각 질의에 대해 조건을 만족하면 YES, 만족하지 않으면 NO를 한 줄에 하나씩 출력한다.

예제 입력

3 6
car
cat
dog
word car
word ca
prefix ca
prefix do
word cart
prefix z

예제 출력

YES
NO
YES
YES
NO
NO

예제 코드

C++20

#include <iostream>
#include <string>

using namespace std;

// 최대 100개 단어 × 길이 20 + 루트
int t[2001][26], sz;
bool fin[2001];

void add(const string& s) {
    int u, x;

    u = 0;
    for (char c : s) {
        x = c - 'a';
        // 해당 문자로 이어지는 노드가 없으면 새 노드 생성
        if (t[u][x] == 0) t[u][x] = ++sz;
        u = t[u][x];
    }
    // 접두사 경로와 완성된 단어를 구분하는 끝 표시
    fin[u] = true;
}

int find(const string& s) {
    int u, x;

    u = 0;
    for (char c : s) {
        x = c - 'a';
        // 해당 접두사 경로가 없으므로 검색 종료
        if (t[u][x] == 0) return -1;
        u = t[u][x];
    }
    return u;
}

int main() {
    int n, q, i, u;
    string s, op;
    bool ok;

    cin >> n >> q;
    for (i = 0; i < n; ++i) {
        cin >> s;
        add(s);
    }

    for (i = 0; i < q; ++i) {
        cin >> op >> s;
        u = find(s);

        // 단어 검색은 끝 표시까지 확인
        if (op == "word") ok = u != -1 && fin[u];
        else ok = u != -1;

        if (ok) cout << "YES\n";
        else cout << "NO\n";
    }
}