Add delete to trie.py + tests

data_structures/trie/trie.py

@@ -1,9 +1,8 @@
 """
 A Trie/Prefix Tree is a kind of search tree used to provide quick lookup
 of words/patterns in a set of words. A basic Trie however has O(n^2) space complexity
-making it impractical in practice. It however provides O(max(search_string, length of longest word)) lookup
-time making it an optimal approach when space is not an issue.
-
+making it impractical in practice. It however provides O(max(search_string, length of longest word)) 
+lookup time making it an optimal approach when space is not an issue.
 """
 
 
@@ -12,7 +11,7 @@ def __init__(self):
         self.nodes = dict()  # Mapping from char to TrieNode
         self.is_leaf = False
 
-    def insert_many(self, words: [str]):  # noqa: E999 This syntax is Python 3 only
+    def insert_many(self, words: [str]):
         """
         Inserts a list of words into the Trie
         :param words: list of string words
@@ -21,7 +20,7 @@ def insert_many(self, words: [str]):  # noqa: E999 This syntax is Python 3 only
         for word in words:
             self.insert(word)
 
-    def insert(self, word: str):  # noqa: E999 This syntax is Python 3 only
+    def insert(self, word: str):
         """
         Inserts a word into the Trie
         :param word: word to be inserted
@@ -34,7 +33,7 @@ def insert(self, word: str):  # noqa: E999 This syntax is Python 3 only
             curr = curr.nodes[char]
         curr.is_leaf = True
 
-    def find(self, word: str) -> bool:  # noqa: E999 This syntax is Python 3 only
+    def find(self, word: str) -> bool:
         """
         Tries to find word in a Trie
         :param word: word to look for
@@ -47,29 +46,82 @@ def find(self, word: str) -> bool:  # noqa: E999 This syntax is Python 3 only
             curr = curr.nodes[char]
         return curr.is_leaf
 
+    def delete(self, word: str):
+        """
+        Deletes a word in a Trie
+        :param word: word to delete
+        :return: None
+        """
 
-def print_words(node: TrieNode, word: str):  # noqa: E999 This syntax is Python 3 only
+        def _delete(curr: TrieNode, word: str, index: int):
+            if index == len(word):
+                # If word does not exist
+                if not curr.is_leaf:
+                    return False
+                curr.is_leaf = False
+                return len(curr.nodes) == 0
+            char = word[index]
+            char_node = curr.nodes.get(char)
+            # If char not in current trie node
+            if not char_node:
+                return False
+            # Flag to check if node can be deleted
+            delete_curr = _delete(char_node, word, index + 1)
+            if delete_curr:
+                del curr.nodes[char]
+                return len(curr.nodes) == 0
+            return delete_curr
+
+        _delete(self, word, 0)
+
+
+def print_words(node: TrieNode, word: str):
     """
     Prints all the words in a Trie
     :param node: root node of Trie
     :param word: Word variable should be empty at start
     :return: None
     """
     if node.is_leaf:
-        print(word, end=' ')
+        print(word, end=" ")
 
     for key, value in node.nodes.items():
         print_words(value, word + key)
 
 
-def test():
-    words = ['banana', 'bananas', 'bandana', 'band', 'apple', 'all', 'beast']
+def test_trie():
+    words = "banana bananas bandana band apple all beast".split()
     root = TrieNode()
     root.insert_many(words)
-    # print_words(root, '')
-    assert root.find('banana')
-    assert not root.find('bandanas')
-    assert not root.find('apps')
-    assert root.find('apple')
+    # print_words(root, "")
+    assert all(root.find(word) for word in words)
+    assert root.find("banana")
+    assert not root.find("bandanas")
+    assert not root.find("apps")
+    assert root.find("apple")
+    assert root.find("all")
+    root.delete("all")
+    assert not root.find("all")
+    root.delete("banana")
+    assert not root.find("banana")
+    assert root.find("bananas")
+    return True
+
+
+def print_results(msg: str, passes: bool) -> None:
+    print(str(msg), "works!" if passes else "doesn't work :(")
+
+
+def pytests():
+    assert test_trie()
+
+
+def main():
+    """
+    >>> pytests()
+    """
+    print_results("Testing trie functionality", test_trie())
+
 
-test()
+if __name__ == "__main__":
+    main()