Auto-complete using Tries in Python with OOP

Question

I am working on a problem to develop auto-complete functionality to practice different applications of Tries. The auto-complete function will return all the possible words in the wordlist given a prefix. I came up with the following solution which returns the result as expected. I need some help to review my code, in order to enhance the OOP structure and improve the code quality over all.

Please let me know your thoughts, any feedback is welcome.

class Trie():
    def __init__(self):
        self.children = {}
        self.flag = False # Flag to represent that a word ends at this node

    def insert(self, word):
        for char in word:
            if char not in self.children:
               self.children[char]= Trie()

            self = self.children[char]
        self.flag = True

    def all_suffixes(self, prefix):
        results = set()

        if self.flag:
            results.add(prefix)

        if not self.children: 
            return results

        return reduce(lambda a, b: a | b,  [node.all_suffixes(prefix + char) for (char, node) in self.children.items()]) 

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

if __name__=="__main__":
    word_list = ['aardvark','ark', 'altimeter','altitude', 'apotactic', 'bagonet', 'boatlip', 
    'carburant', 'chyliferous', 'consonance', 'cyclospondylic', 
    'dictyostele', 'echelon', 'estadal', 'flaunty', 'gesneriaceous', 
    'hygienic', 'infracentral', 'jipijapa', 'lipoceratous', 'melanthaceae']

    c = Trie()

    for word in word_list:
        c.insert(word)
    print c.autocomplete('a')

Output - ['aardvark', 'ark', 'altimeter', 'altitude', 'apotactic']

I was thinking on the lines of having two classes one as Trie() and the other as Autocomplete(). The __init__(self,word_list) for Autocomplete() can then have the for loop code snippet to create the Trie from the word list.

Answer 1

Bug

Put 'a' in word_list: it is not found by autocomplete('a'). This is because all_suffixes doesn't use the result variable when building the result in case the node has children.

Naming

• The name self.flag tells nothing about the purpose of the variable. self.word_end would be better.
• all_suffixes returns whole words, not suffixes. words_with_prefix would be a better name. Alternatively, the function could return suffixes only and adding the prefix would be the responsibility of autocomplete.

Other

Instead of using a set and reduce in all_suffixes, it would be simpler to use just a list and extend it in a loop. (There should be no duplicate words in the list anyway, because the trie does not store duplicates)

def words_with_prefix(self, prefix):
    results = []
    if self.word_end:
        results.append(prefix)
    for (char, node) in self.children.items():
        results.extend(node.all_suffixes(prefix + char))
    return results