[문자열] Trie

작성일 | In

해결해야 할 문제

Trie라는 문자열을 다루는 구조를 설명하기 전에, 아래 문제에 대한 효율적인 풀이를 생각해봅시다.

N개의 정수로 이루어진 집합 A가 있고, M개의 쿼리가 주어집니다. 쿼리의 내용은 아래와 같습니다. N은 -10^18 ~ 10^18 범위에 존재하는 정수입니다.
exist x : 집합 A에 x가 존재하면 true, 반대의 경우에는 false 반환

N의 범위가 작다면 Hash Table을 이용하면 되지만, 이런 경우에는 흔히 set을 구현할 때 사용하는 BBST(Balanced Binary Tree)를 사용하는 것이 효율적입니다. 각 쿼리를 O(log N)만에 처리할 수 있습니다.

Trie 구조

문자열로 넘어와봅시다.
숫자는 비교할 때 O(1)밖에 걸리지 않습니다. 그러나 문자열은 최악의 경우 길이에 비례하게, 즉 O(n)이 걸립니다. 문자열이 N개 있는 BBST에서 탐색 쿼리는 O(길이 * log N)이 걸리게 됩니다.

문자열에 대해 위와 같은 쿼리를 보다 빠르게 처리하기 위해 Trie라는 구조를 사용합니다.

[“A”, “I”, “IT”, “SUN”, “SUA”, “TO”, “TED”, “TEN”]이 주어진다고 가정하면 아래와 같이 트라이를 구성할 수 있습니다.

Trie는 위 그림과 같이 트리를 구성하되, 각 정점이 어떤 문자열의 prefix(접두어)가 되게 합니다. 현재 정점에서 A 간선을 타고 가면 다음 정점에는 뒤에 A가 붙는 식으로 구성이 됩니다.

탐색은 루트에서 시작해 찾고자 하는 정점에 갈 때까지 알맞은 간선을 타고 가면 됩니다.

구현 방법

구현을 해봅시다.
연결 리스트와 유사한 방식으로 구현하는 방법과 vector(동적 배열)를 사용하여 구현하는 방법이 있습니다. 여기에서는 vector를 사용하겠습니다. 또한 여기에서 구현할 Trie는 영어 대문자로 이루어진 단어만 다룰 수 있습니다.
정점을 나타낼 구조체를 정의합시다.

1
2
3
4
5
6
7
8
9
class TrieNode{
	public:
		bool valid;
		int child[26];
		TrieNode(){
			valid = false;
			for(int i=0; i<26; i++) child[i] = -1;
		}
};


이 그림만 봐서는 SU라는 단어가 실제로 존재하는지, 혹은 SUA나 SUN까지 도달하기 위해 거쳐갈 용도로 생성된건지 알 수 없기 때문에 valid라는 변수를 추가해줍니다. valid가 true인 경우, 해당 정점에 있는 단어는 실제로 존재하는 것을 의미합니다.
child배열은 자식 정점이 vector의 몇 번 인덱스에 있는지 저장하는 배열입니다. child[0]에는 뒤에 A가 붙는 자식 정점이, child[25]에는 뒤에 Z가 붙는 자식 정점이 위치한 인덱스가 저장됩니다.

Trie 클래스의 기본적인 골격을 잡아봅시다.

1
2
3
4
5
6
7
8
9
10
11
12
13
class Trie{
  private:
    vector<TrieNode> trie;
    int _newNode(); //새로운 정점 생성 & 해당 정점의 인덱스 반환
    void _add(string &str, int node, int idx); //문자열 삽입
    bool _exist(string &str); //존재하는가?
  public:
    Trie(); //생성자
    void add(String &str); //문자열 삽입(C++ style string)
    void add(char str[]); //문자열 삽입(C style string)
    bool exist(string &str); //문자열 삽입(C++ style string)
    bool exist(char str[]); //문자열 삽입(C style string)
}

private 메소드를 하나씩 구현해봅시다.

_newNode 메소드는 TrieNode를 하나 생성해 trie라는 vector에 넣어주고, 인덱스를 반환하면 됩니다.

1
2
3
4
5
int _newNode(){
  TrieNode tmp;
  trie.push_back(tmp);
  return trie.size() - 1;
}

_add 메소드는 재귀적으로 구현됩니다.
idx는 삽입할 문자열에서 현재 처리하고 있는 인덱스, node는 처리하고 있는 TrieNode의 인덱스입니다. 코드를 보는 것이 이해가 빠를 것이라 생각됩니다.

1
2
3
4
5
6
7
8
9
10
11
void _add(string &str, int node, int idx){
  if(idx == str.size()){
    trie[node].valid = true; return;
  }
  int c = str[idx] - 'A';
  if(trie[node].child[c] == -1){
    int next = _newNode();
    trie[node].child[c] = next;
  }
  _add(str, trie[node].child[c], idx+1);
}

_exist 메소드는 문자열의 각 문자를 순회하면서 알맞은 간선을 타고 갑니다. 타고갈 간선이 없으면 false를 반환하고, 마지막에 도착한 정점의 valid가 false인 경우에도 false를 반환합니다.

1
2
3
4
5
6
7
8
9
bool _exist(string &str){
  int now = 0;
  for(int i=0; i<str.size(); i++){
    int c = str[i] - 'A';
    if(trie[now].child[c] == -1) return false;
    now = trie[now].child[c];
  }
  return trie[now].valid;
}

전체 코드

Trie 클래스의 구현과 사용 예제입니다. (이 링크)에서도 확인 가능합니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
#include <bits/stdc++.h>
using namespace std;

class TrieNode{
	public:
		bool valid;
		int child[26];
		TrieNode(){
			valid = false;
			for(int i=0; i<26; i++) child[i] = -1;
		}
};

class Trie{
	private:
		vector<TrieNode> trie;
		int _newNode(){
			TrieNode tmp;
			trie.push_back(tmp);
			return trie.size() - 1;
		}
		void _add(string &str, int node, int idx){
			if(idx == str.size()){
				trie[node].valid = true; return;
			}
			int c = str[idx] - 'A';
			if(trie[node].child[c] == -1){
				int next = _newNode();
				trie[node].child[c] = next;
			}
			_add(str, trie[node].child[c], idx+1);
		}
		bool _exist(string &str){
			int now = 0;
			for(int i=0; i<str.size(); i++){
				int c = str[i] - 'A';
				if(trie[now].child[c] == -1) return false;
				now = trie[now].child[c];
			}
			return trie[now].valid;
		}
	public:
		Trie(){
			_newNode();
		}
		void add(string &str){
			_add(str, 0, 0);
		}
		void add(char str[]){
			string tmp(str);
			_add(tmp, 0, 0);
		}
		bool exist(string &str){
			return _exist(str);
		}
		bool exist(char str[]){
			string tmp(str);
			return _exist(tmp);
		}
};

int main(){
	Trie a;
	a.add("ASDF");
	a.add("FDSA");
	a.add("NJH");
	a.add("NJHAF");
	printf("%d\n", a.exist("ASDF"));
	printf("%d\n", a.exist("ASD"));
	printf("%d\n", a.exist(""));
	printf("%d\n", a.exist("NJ"));
	printf("%d\n", a.exist("NJH"));
	printf("%d\n", a.exist("NJHA"));
	printf("%d\n", a.exist("NJHAF"));
	printf("%d\n", a.exist("NJHAD"));
}