-
Notifications
You must be signed in to change notification settings - Fork 0
/
trie-search-tree.py
137 lines (98 loc) · 3.72 KB
/
trie-search-tree.py
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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import re
class Node:
def __init__(self):
self.keys = {}
self.end = False
class Trie:
"""
A trie search tree that stores words in an ordered tree structure in which each node
represents an individual letter.
"""
def __init__(self):
self.root = Node()
def add(self, word):
"""Adds a new word to the trie."""
if re.search("^[A-Za-z]+$", word) is None:
raise ValueError(
"Argument must be a single word with no special characters."
)
def add_from(node, remaining_ltrs):
if remaining_ltrs:
# Move to next letter, creating new node if necessary
add_from(
node.keys.setdefault(remaining_ltrs[0], Node()), remaining_ltrs[1:]
)
else:
# Mark end of word
node.end = True
add_from(self.root, word)
def contains(self, string):
"""
Checks if a specific word is currently stored in the trie. Using the `in`
keyword is equivalent to this method (e.g., `string in trie` is the same as
`trie.contains(string)`).
"""
if re.search("^[A-Za-z]+$", string) is None:
return False
def search_from(node, remaining_ltrs):
if remaining_ltrs:
try:
# Move to next letter
return search_from(node.keys[remaining_ltrs[0]], remaining_ltrs[1:])
except:
# String not present
return False
else:
# Word exists if this node is an end point
return node.end
return search_from(self.root, string)
__contains__ = contains
def remove(self, word):
"""
Removes a word from the trie and returns a boolean confirming whether or not the
word was present (and was therefore removed).
"""
if re.search("^[A-Za-z]+$", word) is None:
return False
# List of (current node, next letter) tuples
word_nodes = []
def collect_nodes_from(node, remaining_ltrs):
if remaining_ltrs:
word_nodes.append((node, remaining_ltrs[0]))
try:
# Move to next letter
return collect_nodes_from(
node.keys[remaining_ltrs[0]], remaining_ltrs[1:]
)
except:
# String not present
return False
else:
word_nodes.append((node, None))
# Word exists if this node is an end point
return node.end
if not collect_nodes_from(self.root, word):
return False
# Remove endpoint from last node
current_node, next_ltr = word_nodes.pop()
current_node.end = False
on_shared_path = len(current_node.keys) > 0
# Remove nodes that aren't part of any other word
while word_nodes and not on_shared_path:
current_node, next_ltr = word_nodes.pop()
del current_node.keys[next_ltr]
on_shared_path = len(current_node.keys) > 0 or current_node.end
return True
def clear(self):
"""Clears the contents of the trie."""
self.root = Node()
def get_all(self):
"""Returns a set of all words stored in the trie."""
words = set()
def search_from(node, word_so_far):
if node.end:
words.add(word_so_far)
for letter, next_node in node.keys.items():
search_from(next_node, word_so_far + letter)
search_from(self.root, "")
return words