In order to return the trie
, you may want to store the root
somewhere safe and return that:
def trie_setup_for_words(words):
root = {}
trie = root
for current_word in words:
for current_letter in current_word:
if current_letter not in trie:
trie[current_letter] = {}
trie = trie[current_letter]
# Mark the end of a word
trie['*'] = '*'
# The root is also needed to reset it when a new word is added
trie = root
# Return the root here, not just the last leaf
return root
if __name__ == "__main__":
trie = trie_setup_for_words(["Patrick", 'Pete', 'Peter'])
If we print out the trie
, we will see something like:
{'P': {'a': {'t': {'r': {'i': {'c': {'k': {'*': '*'}}}}}}, 'e': {'t': {'e': {'*': '*', 'r': {'*': '*'}}}}}}
Just to make sure out trie is constructed correctly, we may want to test out if a word exits in the trie
:
def contains_word(trie: dict, word: str):
if not word and '*' in trie:
return True
if word and word[0] in trie:
return contains_word(trie[word[0]], word[1:])
return False
This will return the following results:
print(contains_word(trie, "Patrick")) # Prints True
print(contains_word(trie, "Patric")) # Prints False
print(contains_word(trie, "Pete")) # Prints True
print(contains_word(trie, "Peter")) # Prints True
print(contains_word(trie, "Pet")) # Prints False