Problem #11 [Medium]

This problem was asked by Twitter.

Implement an autocomplete system. That is, given a query string s and a set of all possible query strings, return all strings in the set that have s as a prefix.

For example, given the query string de and the set of strings [dog, deer, deal], return [deer, deal].

Hint: Try preprocessing the dictionary into a more efficient data structure to speed up queries.
经典用Trie树。

class TrieNode:
    def __init__(self) -> None:
        self.children ={}
        self.is_end_of_word =False

class AutocompleteSystem:
    def __init__(self, words) -> None:
        self.root = TrieNode()
        self.build_trie(words)
    
    def build_trie(self, words):
        for word in words:
            node = self.root
            for char in word:
                if char not in node.children:
                    node.children[char] = TrieNode()
                node = node.children[char]
            node.is_end_of_word = True

    def autocomplete(self, prefix):
        node = self.root
        result = []
        for char in prefix:
            if char not in node.children:
                return result
            node = node.children[char]

        self.find_words_with_prefix(node, prefix, result)
        return result

    def find_words_with_prefix(self, node, current_word, result):
        if node.is_end_of_word:
            result.append(current_word)

        for char, child_node in node.children.items():
            self.find_words_with_prefix(child_node, current_word + char, result)


# Create an AutocompleteSystem with a set of words
autocomplete_system = AutocompleteSystem(["dog", "deer", "deal"])

# Query for autocomplete suggestions
result = autocomplete_system.autocomplete("de")
print(result)  # Output: ['deer', 'deal']