Code Sample for juergen.rocks

simple_trie.py

root_node = None


class TrieNode(object):

    def __init__(self):
        self.__children = {}
        self.__result = None

    def __add(self, word: str, rest: str) -> None:
        if len(rest) > 0:
            letter = rest[0].upper()
            if letter not in self.__children.keys():
                self.__children[letter] = TrieNode()
            self.__children[letter].__add(word, rest[1:])
        else:
            self.__result = word.upper()

    def add(self, word: str) -> None:
        self.__add(word, word)

    def search(self, word: str) -> list:
        results = []
        if len(word) > 0:
            letter = word[0].upper()
            if letter in self.__children.keys():
                results.extend(self.__children[letter].search(word[1:]))
        else:
            if self.__result is not None:
                results.append(self.__result)
            [results.extend(self.__children[key].search('')) for key in self.__children.keys()]
        return results


def index_and_search(words: list, search: str) -> list:
    """
    Indiziere Suchbegriffe und suche anschließend darin

    :param list[str] words: Begriffe, aus denen ein Trie aufgebaut werden soll
    :param str search: Begriff, nach dem gesucht werden soll
    :return: Liste mit Treffern
    :rtype: list[str]

    Tests:

    Aus dem Beispiel:
    >>> set(index_and_search(['Hallo', 'Haus', 'Welt', 'Wasser'], 'hau')) == {'HAUS'}
    True

    Mehrere Treffer:
    >>> set(index_and_search(['Hallo', 'Haus', 'Welt', 'Wasser', 'HAL 8999'], 'hal')) == {'HALLO', 'HAL 8999'}
    True

    Keine Treffer:
    >>> index_and_search(['Hallo', 'Haus', 'Welt', 'Wasser', 'HAL 8999'], 'sonne')
    []
    """

    global root_node
    root_node = TrieNode()
    [root_node.add(w) for w in words]
    return root_node.search(search)