Link: https://leetcode.com/problems/implement-trie-prefix-tree
Introduction
Trie 자료구조는 계속 정리해둬야지 하고 있었는데 마침 Leetcode에서 좋은 문제를 발견했다.
arr = [‘A’, ‘to’, ‘tea’, ‘ted’, ‘ten’, ‘t’, ‘in’, ‘inn’]를 Trie에 저장한 상태
Trie는 트리 형태의 자료 구조로, Root Node에서 이어지며 문자열의 글자를 하나씩 저장한다. 저장 공간을 많이 차지한다는 단점을 가지지만 문자열 검색에 탁월한 성능을 가진다는 장점이 있다.
- Insert: O(n)
- 길이가 n인 문자열 삽입
- Search: O(n)
- 길이가 n인 문자열이 자료구조에 포함되어 있는지 검색
뿐만 아니라, Autocomplete, Spell Check등에도 사용할 수 있다.
I. Trie Class
Trie는 노드 클래스를 만들어 연결하는 방식으로 구현해도 되지만, 여기서는 Linked Nodes의 개념만 차용하여 간결하게 Dictionary(Map)로 구현한다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Trie:
def __init__(self):
self.root = {}
def insert(self, word):
pass
def search(self, word):
pass
def startsWith(self, prefix):
pass
def suggest(self, word):
pass
II. Insert
word에서 한 글자 한 글자씩 떼서 Trie에 기록하기 위해 순회한다.
1
2
3
4
5
6
7
class Trie:
def __init__(self):
self.root = {}
def insert(self, word):
for ch in word: # 2-1
pass
Root Node부터 시작한다. 현재 노드를 가리키는 cur_node에 Root Node를 할당한다.
1
2
3
4
5
6
7
8
class Trie:
def __init__(self):
self.root = {}
def insert(self, word):
cur_node = self.root # 2-2
for ch in word:
pass
비어 있는 Trie에 trip이라는 단어를 집어넣는다고 생각해보자.
여기서는 Dictionary를 사용하므로 Key가 해당 노드의 알파벳이고, Value가 해당 노드에 이어져 있는 다른 노드들의 집합이 된다.
‘t’, ‘r’, ‘i’, ‘p’라는 알파벳을 갖고 있는 노드들을 만들어서 이어주어야 하는데, 처음이라 비어있으면 해당 노드를 먼저 만들어줘야한다. 여기서는 해당 노드의 알파벳을 Key로 가지는 집합을 만들어주면 된다.
1
2
3
4
5
6
7
8
9
class Trie:
def __init__(self):
self.root = {}
def insert(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
cur_node[ch] = {} # 2-3
json 형태로 살펴보면 아래와 같은 느낌이다.
1
2
3
{
"t": {}
}
cur_node에서 cur_node[ch]로 넘어가기 위해 cur_node를 cur_node[ch]로 바꿔주자. (쉽게 생각해서 depth를 늘려서 계속 안쪽으로 들어간다는 뜻)
1
2
3
4
5
6
7
8
9
10
class Trie:
def __init__(self):
self.root = {}
def insert(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
cur_node[ch] = {}
cur_node = cur_node[ch] # 2-4
마지막 알파벳인 p가 쓰여져 있는 노드에는 해당 위치에 단어가 있다는 '*' 표시를 남기고 해당 단어를 넣어둔다.
1
2
3
4
5
6
7
8
9
10
11
class Trie:
def __init__(self):
self.root = {}
def insert(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
cur_node[ch] = {}
cur_node = cur_node[ch]
cur_node['*'] = word # 2-5
trip, trap라는 단어들을 넣었을 때 trie는 이런 형태가 된다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
{
"t": {
"r": {
"i": {
"p": {
"*": "trip"
}
},
"a": {
"p": {
"*": "trap"
}
}
}
}
}
III. Search
노드를 계속해서 탐색해가다가 중간에 끊어지면 해당 단어가 Trie에 없다는 뜻이므로 False를 반환한다. 나머지는 Insert랑 똑같이 계속해서 탐색할 노드를 바꿔가며 안으로 들어간다.
1
2
3
4
5
6
7
8
9
10
class Trie:
def __init__(self):
self.root = {}
def search(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node: # 3-1
return False
cur_node = cur_node[ch]
word의 글자 수만큼 들어갔을 때까지 끊기지 않았다면 종착지에 Insert 때 해둔 '*' 표시가 해당 노드에 있는지를 판별하고 있으면 True를, 없다면 False를 반환한다.
1
2
3
4
5
6
7
8
9
10
11
class Trie:
def __init__(self):
self.root = {}
def search(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
return False
cur_node = cur_node[ch]
return '*' in cur_node # 3-2
IV. Prefix
해당 단어로 시작하는 단어가 Trie에 있는지 찾는 법은 search와 크게 다르지 않다. 마지막 노드에 '*' 표시가 없어도 된다는 것만 다르다.
1
2
3
4
5
6
7
8
9
10
11
class Trie:
def __init__(self):
self.root = {}
def startsWith(self, prefix: str) -> bool:
cur_node = self.root
for ch in prefix:
if ch not in cur_node:
return False
cur_node = cur_node[ch]
return True # 4-1
V. Autocomplete
Autocomplete도 입력받은 word의 끝까지 탐색하는 것까지는 같다.
1
2
3
4
5
6
7
8
9
10
class Trie:
def __init__(self):
self.root = {}
def suggest(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
return False
cur_node = cur_node[ch]
여기서 autocomplete에 해당하는 단어들을 담을 result 리스트를 준비해주고 해당 노드 부터 하위 노드들을 모두 탐색해서 단어들을 다 담아 반환해 주면 된다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Trie:
def __init__(self):
self.root = {}
def suggest(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
return False
cur_node = cur_node[ch]
result = []
self.find_all(cur_node, result)
return result
find_all()를 구현해보자. 우선 현재 노드가 가지고 있는 알파벳(key)과 다음 노드들의 집합(value)을 차례로 순회한다.
1
2
3
4
5
6
7
class Trie:
def __init__(self):
self.root = {}
def find_all(self, node, result):
for key, value in node.items(): # 5-1
pass
key가 '*'인 경우에는 현재 노드에 후보 단어가 들어 있으므로 result에 담고 넘어간다.
1
2
3
4
5
6
7
8
9
10
class Trie:
def __init__(self):
self.root = {}
def find_all(self, node, result):
for key, value in node.items():
if key == '*': # 5-2
result.append(value)
continue
그렇지 않은 경우는 해당 노드를 탐색한다.
1
2
3
4
5
6
7
8
9
10
class Trie:
def __init__(self):
self.root = {}
def find_all(self, node, result):
for key, value in node.items():
if key == '*':
result.append(value)
continue
self.find_all(node[key], result) # 5-3
아래는 arr = [“test”, “tennis”, “tri”, “trie”, “trip”, “trap”, “time”]을 넣은 Trie의 모습이다.
1
2
3
4
5
6
7
8
9
10
11
12
13
{
"t": {
"e": {
"s": { "t": { "*": "test" } },
"n": { "n": { "i": { "s": { "*": "tennis" } } } }
},
"r": {
"i": { "e": { "*": "trie" }, "*": "tri", "p": { "*": "trip" } },
"a": { "p": { "*": "trap" } }
},
"i": { "m": { "e": { "*": "time" } } }
}
}
그리고 “tr”이라는 문자열을 가지고 autocomplete를 실행하면 ['trie', 'tri', 'trip', 'trap'] 가 반환된다.
VI. Trie Class (Complete)
중복되는 탐색 부분을 find_target_node()로 묶었다.
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
class Trie:
def __init__(self):
self.root = {}
def find_target_node(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
return None
cur_node = cur_node[ch]
return cur_node
def find_all(self, node, result):
for key, value in node.items():
if key == '*':
result.append(value)
continue
self.find_all(node[key], result)
def insert(self, word):
cur_node = self.root
for ch in word:
if ch not in cur_node:
cur_node[ch] = {}
cur_node = cur_node[ch]
cur_node['*'] = word
def search(self, word):
target_node = self.find_target_node(word)
return target_node is not None and '*' in target_node
def startsWith(self, prefix: str) -> bool:
target_node = self.find_target_node(prefix)
return target_node is not None
def suggest(self, word):
target_node = self.find_target_node(word)
result = []
self.find_all(target_node, result)
return result
References
- https://en.wikipedia.org/wiki/Trie