Trie
"""
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.
"""
class TrieNode:
def __init__(self) -> None:
self.nodes: dict[str, TrieNode] = {} # Mapping from char to TrieNode
self.is_leaf = False
def insert_many(self, words: list[str]) -> None:
"""
Inserts a list of words into the Trie
:param words: list of string words
:return: None
"""
for word in words:
self.insert(word)
def insert(self, word: str) -> None:
"""
Inserts a word into the Trie
:param word: word to be inserted
:return: None
"""
curr = self
for char in word:
if char not in curr.nodes:
curr.nodes[char] = TrieNode()
curr = curr.nodes[char]
curr.is_leaf = True
def find(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
for char in word:
if char not in curr.nodes:
return False
curr = curr.nodes[char]
return curr.is_leaf
def delete(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:
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) -> 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
"""
if node.is_leaf:
print(word, end=" ")
for key, value in node.nodes.items():
print_words(value, word + key)
def test_trie() -> bool:
words = "banana bananas bandana band apple all beast".split()
root = TrieNode()
root.insert_many(words)
# 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() -> None:
assert test_trie()
def main() -> None:
"""
>>> pytests()
"""
print_results("Testing trie functionality", test_trie())
if __name__ == "__main__":
main()
About this Algorithm
A trie (also called a prefix tree) is a tree data structure that shows order, linking parents to children. It is an efficient way of storing objects that have commonalities. A good example would be in storing phone numbers, or strings in general
For the strings example, supposing we have a list of strings to store in our data store
- egg
- eat
- ear
- end
And one of the methods we are to support is a search operation for any of the words, we can approach it the basic way - select each word, and do a string comparison, matching letter to letter. The algorithm would be as follows:
## searching for ear in data store
data_store = ["egg", "eat", "ear", "end"]
to_find = "ear"
## pick each word
## do a string match letter by letter
## when you find a mismatch, move to the next string
## continue this process
## if at the end of an iteration, index has been increased to
## the length of the word to find, we have found a match
for word in data_store:
index = 0
while index < len(word):
if to_find[index] != word[index]:
break
index += 1
if index == len(to_find):
print("a match has been found")
Without a doubt, this strategy will work, but the time complexity of doing this is O(num of words x len of longest word) which is quite expensive. However, if we represent the storage of numbers in a tree such that each letter appears only once in a particular level in the tree, we can achieve a much better search time. Take, for example, the tree below
e
/ | \
a n g
/ \ | |
r t d g
You can see from the above representation, that all the words are in the tree, starting from the letter e, which is found at the beginning of all the words, then a, n, and g coming in the next level and so on... The above representation is called a trie.
Standard Trie Operations
- insert(): inserts the string into the trie.
- search(): searches for the string within the trie.
Building a Trie
Defining a node class for the elements of the trie
To start building a trie, you first need to define a node with the revelant attributes needed for any trie.
class Node:
def __init__(self, is_word: bool=False):
self.is_word = is_word
self.children = {}
Here, you can see that the class Node has three instance attributes:
- is_word: bool = to mark whether that node in the trie marks the completion of a word
- children: Dict = to hold pointers to other children nodes
Then the trie gets built by creating a node for each letter and adding it as a child to the node before it
Building the trie itself
Start by initializing an empty node
class Trie:
def __init__(self):
self.node = Node()
For the insert operation, fetch the starting node, then for every letter in the word, add it to the children of the letter before it. The final node has its is_word attribute marked as True because we want to be aware of where the word ends
def insert(self, word: str) -> None:
node = self.node
for ltr in word:
if ltr not in node.children:
node.children[ltr] = Node()
node = node.children[ltr]
node.is_word=True
In the code above, the node variable starts by holding a reference to the null node, while the ltr iterating variable starts by holding the first letter in word. This would ensure that node is one level ahead of ltr. As they are both moved forward in the iterations, node will always remain one level ahead of ltr
For the search operation, fetch the starting node, then for every letter in the word, check if it is present in the children attribute of the current node. As long as it is present, repeat for the next letter and next node. If during the search process, we find a letter that is not present, then the word does not exist in the trie. If we successfully get to the end of the iteration, then we have found what we are looking for. It is time to return a value
Take a look at the code
def search(self, word: str) -> bool:
node = self.node
for ltr in word:
if ltr not in node.children:
return False
node = node.children[ltr]
return node.is_word
For the return value, there are two cases:
- we are searching for a word -> return
node.is_wordbecause we want to be sure it is actually a word, and not a prefix - we are searching for a prefix -> return True because whether it is a word or not, it is prefix that exists in the trie
Now here is the full code
class Node:
def __init__(self, is_word: bool=False):
self.is_word = is_word
self.children = {}
class Trie:
def __init__(self):
self.node = Node()
def insert(self, word: str) -> None:
node = self.node
for ltr in word:
if ltr not in node.children:
node.children[ltr] = Node()
node = node.children[ltr]
node.is_word=True
def search(self, word: str) -> bool:
node = self.node
for ltr in word:
if ltr not in node.children:
return False
node = node.children[ltr]
return node.is_word
