Prefix search by trie tree I'm working on the prefix search problem:


  Given a set of words, for example words = ['a', 'apple', 'angle', 'angel', 'bat', 'bats'], for any given prefix, find all matched words. For example, if input is ang, return 'angle', 'angel', if no match, return an empty list [].


Any advice on performance improvement in terms of algorithm time complexity (not sure if Trie Tree is the best solution), code bugs or general advice on code style is highly appreciated.

from collections import defaultdict
class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.isEnd = False
    def insert(self, word):
        node = self
        for w in word:
            node = node.children[w]
        node.isEnd = True
    def search(self, word):
        node = self
        for w in word:
            if w in node.children:
                node = node.children[w]
            else:
                return []
        # prefix match
        # traverse currnt node to all leaf nodes
        result = []
        self.traverse(node, list(word), result)
        return [''.join(r) for r in result]

    def traverse(self, root, prefix, result):
        if root.isEnd:
            result.append(prefix[:])
        for c,n in root.children.items():
            prefix.append(c)
            self.traverse(n, prefix, result)
            prefix.pop(-1)
if __name__ == "__main__":
    words = ['a', 'apple', 'angle', 'angel', 'bat', 'bats']
    root = TrieNode()
    for w in words:
        root.insert(w)
    print root.search('a') # 'a', 'apple', 'angle', 'angel'
    print root.search('ang') # 'angle', 'angel'
    print root.search('angl') # 'angle'
    print root.search('z') # []
What comes to the data structure, I don't think it gets any more efficient, not at least as long as you confine yourself to search by prefix. For that reason, I will "review" your actual code and not the algorithm.

self.isEnd = False


I believe more canonic name would be self.is_leaf. Note that isEnd is not idiomatic Python, is_end is.

You should have two empty lines before any class declaration.

import ...


class FooBar:
    pass


Also, in

for c,n in ...


PEP 8 requires a space after (each) comma:

for c, n in ...
      ^
    space!


What comes to renaming self:

def insert(self, word):
    node = self
    for w in word:
        node = node.children[w]
    node.isEnd = True


I would not do it, yet I have to admit that it is just an opinion.