Skip to content

Commit b2165a6

Browse files
acrulopezAlex de la Cruzcclauss
authored
Added Radix Tree in data structures (#6616)
* added radix tree to data structures * added doctests * solved flake8 * added type hints * added description for delete function * Update data_structures/trie/radix_tree.py * Update radix_tree.py * Update radix_tree.py * Update radix_tree.py Co-authored-by: Alex de la Cruz <[email protected]> Co-authored-by: Christian Clauss <[email protected]>
1 parent 0fd1ccb commit b2165a6

File tree

1 file changed

+223
-0
lines changed

1 file changed

+223
-0
lines changed

Diff for: data_structures/trie/radix_tree.py

+223
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
"""
2+
A Radix Tree is a data structure that represents a space-optimized
3+
trie (prefix tree) in whicheach node that is the only child is merged
4+
with its parent [https://en.wikipedia.org/wiki/Radix_tree]
5+
"""
6+
7+
8+
class RadixNode:
9+
def __init__(self, prefix: str = "", is_leaf: bool = False) -> None:
10+
# Mapping from the first character of the prefix of the node
11+
self.nodes: dict[str, RadixNode] = {}
12+
13+
# A node will be a leaf if the tree contains its word
14+
self.is_leaf = is_leaf
15+
16+
self.prefix = prefix
17+
18+
def match(self, word: str) -> tuple[str, str, str]:
19+
"""Compute the common substring of the prefix of the node and a word
20+
21+
Args:
22+
word (str): word to compare
23+
24+
Returns:
25+
(str, str, str): common substring, remaining prefix, remaining word
26+
27+
>>> RadixNode("myprefix").match("mystring")
28+
('my', 'prefix', 'string')
29+
"""
30+
x = 0
31+
for q, w in zip(self.prefix, word):
32+
if q != w:
33+
break
34+
35+
x += 1
36+
37+
return self.prefix[:x], self.prefix[x:], word[x:]
38+
39+
def insert_many(self, words: list[str]) -> None:
40+
"""Insert many words in the tree
41+
42+
Args:
43+
words (list[str]): list of words
44+
45+
>>> RadixNode("myprefix").insert_many(["mystring", "hello"])
46+
"""
47+
for word in words:
48+
self.insert(word)
49+
50+
def insert(self, word: str) -> None:
51+
"""Insert a word into the tree
52+
53+
Args:
54+
word (str): word to insert
55+
56+
>>> RadixNode("myprefix").insert("mystring")
57+
"""
58+
# Case 1: If the word is the prefix of the node
59+
# Solution: We set the current node as leaf
60+
if self.prefix == word:
61+
self.is_leaf = True
62+
63+
# Case 2: The node has no edges that have a prefix to the word
64+
# Solution: We create an edge from the current node to a new one
65+
# containing the word
66+
elif word[0] not in self.nodes:
67+
self.nodes[word[0]] = RadixNode(prefix=word, is_leaf=True)
68+
69+
else:
70+
incoming_node = self.nodes[word[0]]
71+
matching_string, remaining_prefix, remaining_word = incoming_node.match(
72+
word
73+
)
74+
75+
# Case 3: The node prefix is equal to the matching
76+
# Solution: We insert remaining word on the next node
77+
if remaining_prefix == "":
78+
self.nodes[matching_string[0]].insert(remaining_word)
79+
80+
# Case 4: The word is greater equal to the matching
81+
# Solution: Create a node in between both nodes, change
82+
# prefixes and add the new node for the remaining word
83+
else:
84+
incoming_node.prefix = remaining_prefix
85+
86+
aux_node = self.nodes[matching_string[0]]
87+
self.nodes[matching_string[0]] = RadixNode(matching_string, False)
88+
self.nodes[matching_string[0]].nodes[remaining_prefix[0]] = aux_node
89+
90+
if remaining_word == "":
91+
self.nodes[matching_string[0]].is_leaf = True
92+
else:
93+
self.nodes[matching_string[0]].insert(remaining_word)
94+
95+
def find(self, word: str) -> bool:
96+
"""Returns if the word is on the tree
97+
98+
Args:
99+
word (str): word to check
100+
101+
Returns:
102+
bool: True if the word appears on the tree
103+
104+
>>> RadixNode("myprefix").find("mystring")
105+
False
106+
"""
107+
incoming_node = self.nodes.get(word[0], None)
108+
if not incoming_node:
109+
return False
110+
else:
111+
matching_string, remaining_prefix, remaining_word = incoming_node.match(
112+
word
113+
)
114+
# If there is remaining prefix, the word can't be on the tree
115+
if remaining_prefix != "":
116+
return False
117+
# This applies when the word and the prefix are equal
118+
elif remaining_word == "":
119+
return incoming_node.is_leaf
120+
# We have word remaining so we check the next node
121+
else:
122+
return incoming_node.find(remaining_word)
123+
124+
def delete(self, word: str) -> bool:
125+
"""Deletes a word from the tree if it exists
126+
127+
Args:
128+
word (str): word to be deleted
129+
130+
Returns:
131+
bool: True if the word was found and deleted. False if word is not found
132+
133+
>>> RadixNode("myprefix").delete("mystring")
134+
False
135+
"""
136+
incoming_node = self.nodes.get(word[0], None)
137+
if not incoming_node:
138+
return False
139+
else:
140+
matching_string, remaining_prefix, remaining_word = incoming_node.match(
141+
word
142+
)
143+
# If there is remaining prefix, the word can't be on the tree
144+
if remaining_prefix != "":
145+
return False
146+
# We have word remaining so we check the next node
147+
elif remaining_word != "":
148+
return incoming_node.delete(remaining_word)
149+
else:
150+
# If it is not a leaf, we don't have to delete
151+
if not incoming_node.is_leaf:
152+
return False
153+
else:
154+
# We delete the nodes if no edges go from it
155+
if len(incoming_node.nodes) == 0:
156+
del self.nodes[word[0]]
157+
# We merge the current node with its only child
158+
if len(self.nodes) == 1 and not self.is_leaf:
159+
merging_node = list(self.nodes.values())[0]
160+
self.is_leaf = merging_node.is_leaf
161+
self.prefix += merging_node.prefix
162+
self.nodes = merging_node.nodes
163+
# If there is more than 1 edge, we just mark it as non-leaf
164+
elif len(incoming_node.nodes) > 1:
165+
incoming_node.is_leaf = False
166+
# If there is 1 edge, we merge it with its child
167+
else:
168+
merging_node = list(incoming_node.nodes.values())[0]
169+
incoming_node.is_leaf = merging_node.is_leaf
170+
incoming_node.prefix += merging_node.prefix
171+
incoming_node.nodes = merging_node.nodes
172+
173+
return True
174+
175+
def print_tree(self, height: int = 0) -> None:
176+
"""Print the tree
177+
178+
Args:
179+
height (int, optional): Height of the printed node
180+
"""
181+
if self.prefix != "":
182+
print("-" * height, self.prefix, " (leaf)" if self.is_leaf else "")
183+
184+
for value in self.nodes.values():
185+
value.print_tree(height + 1)
186+
187+
188+
def test_trie() -> bool:
189+
words = "banana bananas bandana band apple all beast".split()
190+
root = RadixNode()
191+
root.insert_many(words)
192+
193+
assert all(root.find(word) for word in words)
194+
assert not root.find("bandanas")
195+
assert not root.find("apps")
196+
root.delete("all")
197+
assert not root.find("all")
198+
root.delete("banana")
199+
assert not root.find("banana")
200+
assert root.find("bananas")
201+
202+
return True
203+
204+
205+
def pytests() -> None:
206+
assert test_trie()
207+
208+
209+
def main() -> None:
210+
"""
211+
>>> pytests()
212+
"""
213+
root = RadixNode()
214+
words = "banana bananas bandanas bandana band apple all beast".split()
215+
root.insert_many(words)
216+
217+
print("Words:", words)
218+
print("Tree:")
219+
root.print_tree()
220+
221+
222+
if __name__ == "__main__":
223+
main()

0 commit comments

Comments
 (0)