diff --git a/data_structures/trie/radix_tree.py b/data_structures/trie/radix_tree.py index caf566a6ce30..b227d774858e 100644 --- a/data_structures/trie/radix_tree.py +++ b/data_structures/trie/radix_tree.py @@ -1,9 +1,11 @@ """ 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 +trie (prefix tree) in which each node that is the only child is merged with its parent [https://en.wikipedia.org/wiki/Radix_tree] """ +import unittest + class RadixNode: def __init__(self, prefix: str = "", is_leaf: bool = False) -> None: @@ -62,6 +64,11 @@ def insert(self, word: str) -> None: -- A (leaf) --- A (leaf) """ + ## Handle the Case where the word is empty by using an if branch + if word == "": + self.is_leaf = True + return + # Case 1: If the word is the prefix of the node # Solution: We set the current node as leaf if self.prefix == word and not self.is_leaf: @@ -80,11 +87,11 @@ def insert(self, word: str) -> None: ) # Case 3: The node prefix is equal to the matching - # Solution: We insert remaining word on the next node + # Solution: We insert the remaining word on the next node if remaining_prefix == "": self.nodes[matching_string[0]].insert(remaining_word) - # Case 4: The word is greater equal to the matching + # Case 4: The word is greater than or equal to the matching # Solution: Create a node in between both nodes, change # prefixes and add the new node for the remaining word else: @@ -100,7 +107,7 @@ def insert(self, word: str) -> None: self.nodes[matching_string[0]].insert(remaining_word) def find(self, word: str) -> bool: - """Returns if the word is on the tree + """Returns whether the word is on the tree Args: word (str): word to check @@ -191,7 +198,7 @@ def print_tree(self, height: int = 0) -> None: value.print_tree(height + 1) -def test_trie() -> bool: +def test_trie() -> None: words = "banana bananas bandana band apple all beast".split() root = RadixNode() root.insert_many(words) @@ -205,25 +212,41 @@ def test_trie() -> bool: assert not root.find("banana") assert root.find("bananas") - return True +class TestRadixNode(unittest.TestCase): + def test_trie(self) -> None: + words = "banana bananas bandana band apple all beast".split() + root = RadixNode() + root.insert_many(words) -def pytests() -> None: - assert test_trie() + assert all(root.find(word) for word in words) + assert not root.find("bandanas") + assert not root.find("apps") + root.delete("all") + assert not root.find("all") + root.delete("banana") + assert not root.find("banana") + assert root.find("bananas") + def test_trie_2(self) -> None: + """ + Now add a new test case that inserts + foobbb, fooaaa, foo in the given order and checks + for different assertions + """ + words = "foobbb fooaaa foo".split() + root = RadixNode() + root.insert_many(words) -def main() -> 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() + assert all(root.find(word) for word in words) + root.delete("foo") + assert not root.find("foo") + assert root.find("foobbb") + assert root.find("fooaaa") if __name__ == "__main__": - main() + import doctest + + doctest.testmod() + unittest.main()