trie.py

"""
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.
"""


classTrieNode:
def__init__(self) ->None:
self.nodes: dict[str, TrieNode] = {} # Mapping from char to TrieNode
self.is_leaf=False

definsert_many(self, words: list[str]) ->None:
"""
 Inserts a list of words into the Trie
 :param words: list of string words
 :return: None
 """
forwordinwords:
self.insert(word)

definsert(self, word: str) ->None:
"""
 Inserts a word into the Trie
 :param word: word to be inserted
 :return: None
 """
curr=self
forcharinword:
ifcharnotincurr.nodes:
curr.nodes[char] =TrieNode()
curr=curr.nodes[char]
curr.is_leaf=True

deffind(self, word: str) ->bool:
"""
 Tries to find word in a Trie
 :param word: word to look for
 :return: Returns True if word is found, False otherwise
 """
curr=self
forcharinword:
ifcharnotincurr.nodes:
returnFalse
curr=curr.nodes[char]
returncurr.is_leaf

defdelete(self, word: str) ->None:
"""
 Deletes a word in a Trie
 :param word: word to delete
 :return: None
 """

def_delete(curr: TrieNode, word: str, index: int) ->bool:
ifindex==len(word):
# If word does not exist
ifnotcurr.is_leaf:
returnFalse
curr.is_leaf=False
returnlen(curr.nodes) ==0
char=word[index]
char_node=curr.nodes.get(char)
# If char not in current trie node
ifnotchar_node:
returnFalse
# Flag to check if node can be deleted
delete_curr=_delete(char_node, word, index+1)
ifdelete_curr:
delcurr.nodes[char]
returnlen(curr.nodes) ==0
returndelete_curr

_delete(self, word, 0)


defprint_words(node: TrieNode, word: str) ->None:
"""
 Prints all the words in a Trie
 :param node: root node of Trie
 :param word: Word variable should be empty at start
 :return: None
 """
ifnode.is_leaf:
print(word, end=" ")

forkey, valueinnode.nodes.items():
print_words(value, word+key)


deftest_trie() ->bool:
words="banana bananas bandana band apple all beast".split()
root=TrieNode()
root.insert_many(words)
# print_words(root, "")
assertall(root.find(word) forwordinwords)
assertroot.find("banana")
assertnotroot.find("bandanas")
assertnotroot.find("apps")
assertroot.find("apple")
assertroot.find("all")
root.delete("all")
assertnotroot.find("all")
root.delete("banana")
assertnotroot.find("banana")
assertroot.find("bananas")
returnTrue


defprint_results(msg: str, passes: bool) ->None:
print(str(msg), "works!"ifpasseselse"doesn't work :(")


defpytests() ->None:
asserttest_trie()


defmain() ->None:
"""
 >>> pytests()
 """
print_results("Testing trie functionality", test_trie())


if__name__=="__main__":
main()