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 = "bananabananasbandanabandappleallbeast".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()