radix_tree.py

"""
A Radix Tree is a data structure that represents a space-optimized
trie (prefix tree) in whicheach node that is the only child is merged
with its parent [https://en.wikipedia.org/wiki/Radix_tree]
"""


classRadixNode:
def__init__(self, prefix: str="", is_leaf: bool=False) ->None:
# Mapping from the first character of the prefix of the node
self.nodes: dict[str, RadixNode] = {}

# A node will be a leaf if the tree contains its word
self.is_leaf=is_leaf

self.prefix=prefix

defmatch(self, word: str) ->tuple[str, str, str]:
"""Compute the common substring of the prefix of the node and a word

 Args:
 word (str): word to compare

 Returns:
 (str, str, str): common substring, remaining prefix, remaining word

 >>> RadixNode("myprefix").match("mystring")
 ('my', 'prefix', 'string')
 """
x=0
forq, winzip(self.prefix, word):
ifq!=w:
break

x+=1

returnself.prefix[:x], self.prefix[x:], word[x:]

definsert_many(self, words: list[str]) ->None:
"""Insert many words in the tree

 Args:
 words (list[str]): list of words

 >>> RadixNode("myprefix").insert_many(["mystring", "hello"])
 """
forwordinwords:
self.insert(word)

definsert(self, word: str) ->None:
"""Insert a word into the tree

 Args:
 word (str): word to insert

 >>> RadixNode("myprefix").insert("mystring")

 >>> root = RadixNode()
 >>> root.insert_many(['myprefix', 'myprefixA', 'myprefixAA'])
 >>> root.print_tree()
 - myprefix (leaf)
 -- A (leaf)
 --- A (leaf)
 """
# Case 1: If the word is the prefix of the node
# Solution: We set the current node as leaf
ifself.prefix==wordandnotself.is_leaf:
self.is_leaf=True

# Case 2: The node has no edges that have a prefix to the word
# Solution: We create an edge from the current node to a new one
# containing the word
elifword[0] notinself.nodes:
self.nodes[word[0]] =RadixNode(prefix=word, is_leaf=True)

else:
incoming_node=self.nodes[word[0]]
matching_string, remaining_prefix, remaining_word=incoming_node.match(
word
 )

# Case 3: The node prefix is equal to the matching
# Solution: We insert remaining word on the next node
ifremaining_prefix=="":
self.nodes[matching_string[0]].insert(remaining_word)

# Case 4: The word is greater equal to the matching
# Solution: Create a node in between both nodes, change
# prefixes and add the new node for the remaining word
else:
incoming_node.prefix=remaining_prefix

aux_node=self.nodes[matching_string[0]]
self.nodes[matching_string[0]] =RadixNode(matching_string, False)
self.nodes[matching_string[0]].nodes[remaining_prefix[0]] =aux_node

ifremaining_word=="":
self.nodes[matching_string[0]].is_leaf=True
else:
self.nodes[matching_string[0]].insert(remaining_word)

deffind(self, word: str) ->bool:
"""Returns if the word is on the tree

 Args:
 word (str): word to check

 Returns:
 bool: True if the word appears on the tree

 >>> RadixNode("myprefix").find("mystring")
 False
 """
incoming_node=self.nodes.get(word[0], None)
ifnotincoming_node:
returnFalse
else:
matching_string, remaining_prefix, remaining_word=incoming_node.match(
word
 )
# If there is remaining prefix, the word can't be on the tree
ifremaining_prefix!="":
returnFalse
# This applies when the word and the prefix are equal
elifremaining_word=="":
returnincoming_node.is_leaf
# We have word remaining so we check the next node
else:
returnincoming_node.find(remaining_word)

defdelete(self, word: str) ->bool:
"""Deletes a word from the tree if it exists

 Args:
 word (str): word to be deleted

 Returns:
 bool: True if the word was found and deleted. False if word is not found

 >>> RadixNode("myprefix").delete("mystring")
 False
 """
incoming_node=self.nodes.get(word[0], None)
ifnotincoming_node:
returnFalse
else:
matching_string, remaining_prefix, remaining_word=incoming_node.match(
word
 )
# If there is remaining prefix, the word can't be on the tree
ifremaining_prefix!="":
returnFalse
# We have word remaining so we check the next node
elifremaining_word!="":
returnincoming_node.delete(remaining_word)
# If it is not a leaf, we don't have to delete
elifnotincoming_node.is_leaf:
returnFalse
else:
# We delete the nodes if no edges go from it
iflen(incoming_node.nodes) ==0:
delself.nodes[word[0]]
# We merge the current node with its only child
iflen(self.nodes) ==1andnotself.is_leaf:
merging_node=next(iter(self.nodes.values()))
self.is_leaf=merging_node.is_leaf
self.prefix+=merging_node.prefix
self.nodes=merging_node.nodes
# If there is more than 1 edge, we just mark it as non-leaf
eliflen(incoming_node.nodes) >1:
incoming_node.is_leaf=False
# If there is 1 edge, we merge it with its child
else:
merging_node=next(iter(incoming_node.nodes.values()))
incoming_node.is_leaf=merging_node.is_leaf
incoming_node.prefix+=merging_node.prefix
incoming_node.nodes=merging_node.nodes

returnTrue

defprint_tree(self, height: int=0) ->None:
"""Print the tree

 Args:
 height (int, optional): Height of the printed node
 """
ifself.prefix!="":
print("-"*height, self.prefix, " (leaf)"ifself.is_leafelse"")

forvalueinself.nodes.values():
value.print_tree(height+1)


deftest_trie() ->bool:
words="banana bananas bandana band apple all beast".split()
root=RadixNode()
root.insert_many(words)

assertall(root.find(word) forwordinwords)
assertnotroot.find("bandanas")
assertnotroot.find("apps")
root.delete("all")
assertnotroot.find("all")
root.delete("banana")
assertnotroot.find("banana")
assertroot.find("bananas")

returnTrue


defpytests() ->None:
asserttest_trie()


defmain() ->None:
"""
 >>> pytests()
 """
root=RadixNode()
words="banana bananas bandanas bandana band apple all beast".split()
root.insert_many(words)

print("Words:", words)
print("Tree:")
root.print_tree()


if__name__=="__main__":
main()