Trie: Autocomplete - ALawliet/algorithms GitHub Wiki

you are given a prefix and a list of words, return all the words that would match

'do'
['dog','dark','cat','door','dodge']

use trie, get to the prefix O(k), then DFS to return the words O(n)
O(k) + O(n)

class TrieNode:
    def __init__(self, children, end):
        self.children = children
        self.end = end

class Solution:
    def __init__(self):
        self.root = TrieNode({}, False)

    def build(self, words):
        for word in words:
            cur = self.root
            for c in word:
                if c not in cur.children:
                    cur.children[c] = TrieNode({}, False)
                cur = cur.children[c]
            cur.end = True

    def autocomplete(self, prefix):
        cur = self.root
        # first get to the end of the prefix
        for c in prefix:
            if c not in cur.children:
                return []
            cur = cur.children[c]
        # then walk the tree from the end of the prefix
        return self.dfs(cur, prefix)

    def dfs(self, cur, prefix):
        words = []
        if cur.end:
            words += [prefix]
        for c in cur.children:
            words += self.dfs(cur.children[c], prefix + c)
        return words

s = Solution()
s.build(['dog', 'dark', 'cat', 'door', 'dodge'])
print(s.autocomplete('do'))
# ['dog', 'door', 'dodge']