Class Trie

strings/autocomplete_using_trie.py:6–32

class Trie:
    def __init__(self) -> None:
        self._trie: dict = {}

    def insert_word(self, text: str) -> None:
        trie = self._trie
        for char in text:
            if char not in trie:
                trie[char] = {}
            trie = trie[char]
        trie[END] = True

    def find_word(self, prefix: str) -> tuple | list:
        trie = self._trie
        for char in prefix:
            if char in trie:
                trie = trie[char]
            else:
                return []
        return self._elements(trie)

    def _elements(self, d: dict) -> tuple:
        result = []
        for c, v in d.items():
            sub_result = [" "] if c == END else [(c + s) for s in self._elements(v)]
            result.extend(sub_result)
        return tuple(result)


trie = Trie()