Autocomplete Using Trie

from __future__ import annotations

END = "#"


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():
            if c == END:
                sub_result = [" "]
            else:
                sub_result = [c + s for s in self._elements(v)]
            result.extend(sub_result)
        return tuple(result)


trie = Trie()
words = ("depart", "detergent", "daring", "dog", "deer", "deal")
for word in words:
    trie.insert_word(word)


def autocomplete_using_trie(string: str) -> tuple:
    """
    >>> trie = Trie()
    >>> for word in words:
    ...     trie.insert_word(word)
    ...
    >>> matches = autocomplete_using_trie("de")
    >>> "detergent " in matches
    True
    >>> "dog " in matches
    False
    """
    suffixes = trie.find_word(string)
    return tuple(string + word for word in suffixes)


def main() -> None:
    print(autocomplete_using_trie("de"))


if __name__ == "__main__":
    import doctest

    doctest.testmod()
    main()
Algerlogo

Β© Alger 2022

About us

We are a group of programmers helping each other build new things, whether it be writing complex encryption programs, or simple ciphers. Our goal is to work together to document and model beautiful, helpful and interesting algorithms using code. We are an open-source community - anyone can contribute. We check each other's work, communicate and collaborate to solve problems. We strive to be welcoming, respectful, yet make sure that our code follows the latest programming guidelines.