trie.py

#
# Original code by:
# https://wangyy395.medium.com/implement-a-trie-in-python-e8dd5c5fde3a
#

from collections import defaultdict

class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.is_end = False

class Trie:
    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        """
        Inserts a word into the trie.
        """
        current = self.root
        for letter in word:
            current = current.children[letter]
        current.is_end = True

    def search(self, word: str) -> bool:
        """
        Returns if the word is in the trie.
        """
        current = self.root
        for letter in word:
            current = current.children.get(letter)
            if current is None:
                return False
        return current.is_end

    def starts_with(self, prefix: str) -> bool:
        """
        Returns if there is any word in the trie that starts with the given prefix.
        """
        current = self.root

        for letter in prefix:
            current = current.children.get(letter)
            if not current:
                return False

        return True

if __name__ == '__main__':
    t = Trie()
    t.insert('águia')
    t.insert('água')

    palavra = 'água'
    print(f"A palavra {palavra} existe na Trie? {t.search(palavra)}" )

    prefixo = 'águ'
    print(f"Existem palavras que possuem o prefixo '{prefixo}'? {t.starts_with(prefixo)}")