diff --git a/data_structures/trie/radix_tree.py b/data_structures/trie/radix_tree.py index caf566a6ce30..7bae3c154bcf 100644 --- a/data_structures/trie/radix_tree.py +++ b/data_structures/trie/radix_tree.py @@ -1,3 +1,5 @@ +import unittest + """ 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 @@ -62,6 +64,11 @@ def insert(self, word: str) -> None: -- A (leaf) --- A (leaf) """ + ## Handle the Case where 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: @@ -191,39 +198,42 @@ def print_tree(self, height: int = 0) -> None: value.print_tree(height + 1) -def test_trie() -> bool: - words = "banana bananas bandana band apple all beast".split() - root = RadixNode() - root.insert_many(words) - - 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") +## write unit test for the code using unittest library with +## logic similar to test_trie() function +## and call it from main() - 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 which 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() + unittest.main()