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]
                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:
        for c,n in root.children.items():
            self.traverse(n, prefix, result)
if __name__ == "__main__":
    words = ['a', 'apple', 'angle', 'angel', 'bat', 'bats']
    root = TrieNode()
    for w in words:
    print root.search('a') # 'a', 'apple', 'angle', 'angel'
    print root.search('ang') # 'angle', 'angel'
    print root.search('angl') # 'angle'
    print root.search('z') # []
  • 3
    \$\begingroup\$ How do you estimate the ratio of adding Words and searching for prefix? I think this is a huge factor to tell whether tree-search is efficient or not. \$\endgroup\$ – SchreiberLex Feb 14 '17 at 8:47
  • \$\begingroup\$ @Lex, if search is relatively less frequent, what is your suggestion? Nice question and vote up. \$\endgroup\$ – Lin Ma Feb 15 '17 at 7:55
  • 1
    \$\begingroup\$ I can't tell you a way to go since this is python i did not respect, since dicts are a memory-efficient way implementing a tree with few more expensiv accessing these values this seems to be a good middle way. Its also depending on your use-case, since it can be efficient to just store strings after another and brute force them if this operation is not needed frequently and is fine to take more time but memory is limited. But everything always with a "depends" ... \$\endgroup\$ – SchreiberLex Feb 16 '17 at 9:29
  • \$\begingroup\$ @Lex, suppose you can use any programming language like Java, my question is more about if from algorithm/data structure perspective, if Trie tree is the best choice for this problem? \$\endgroup\$ – Lin Ma Feb 16 '17 at 23:18

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:

Also, in

for c,n in ...

PEP 8 requires a space after (each) comma:

for c, n in ...

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.

  • \$\begingroup\$ Thanks for the advice and vote up. If search is by prefix (this is the only search pattern, input is prefix, and output is all matched words), what data structure do you think best fits the problem (if not Trie Tree)? \$\endgroup\$ – Lin Ma Feb 15 '17 at 7:57
  • 1
    \$\begingroup\$ @LinMa Well, I am not much of a string algorithms expert, yet it would seem that your implementation is as efficient as it gets. In case you need to do something more fancy, you might need very advanced data structures such as suffix tries, suffix array, and the like. \$\endgroup\$ – coderodde Feb 15 '17 at 8:09
  • \$\begingroup\$ Thanks, but I think Trie tree is suffix tries. Isn't it? :) \$\endgroup\$ – Lin Ma Feb 15 '17 at 8:12
  • 1
    \$\begingroup\$ @LinMa No, suffix trie is built for one word only and contains all the suffixes of that word. Building suffix tries in linear time is highly non-trivial. \$\endgroup\$ – coderodde Feb 15 '17 at 8:42
  • \$\begingroup\$ Thanks coderodde, how suffix tree related to this problem? I think we need to build Trie Tree for all dictionary words, since we need a full prefix match. Do you think build suffix tree will be faster than Trie tree or? \$\endgroup\$ – Lin Ma Feb 16 '17 at 23:16

The Trie definitely fits the problem greatly. Here are some other points, adding to the @coderodde's awesome answer:

  • I don't particularly like the way you traverse the trie to get all paths to the leaf nodes. I would make traverse() method a generator:

    def traverse(self, root, prefix):
        if root.is_leaf:
            yield prefix
        for c, n in root.children.items():
            yield from self.traverse(n, prefix)  # Python 3.3+

    Then, you can improve your search() method by returning from traverse():

    return [''.join(r) for r in self.traverse(node, list(word))]
  • you can define __slots__ to improve on memory usage and performance:

    class TrieNode:
        __slots__ = ['children', 'is_leaf']
  • note the "currnt" typo

  • put 2 newlines after the import statements (PEP8 reference)
  • 1
    \$\begingroup\$ Thanks alecxe, good advice and vote up. I have kinds of feeling two separate function search and traverse seems a bit messy, do you think there is a way to write one elegant function which combine them? \$\endgroup\$ – Lin Ma Feb 15 '17 at 7:56
  • \$\begingroup\$ Hi alecxe, if you could comment on my above question, it will be great. \$\endgroup\$ – Lin Ma Feb 16 '17 at 23:17
  • 1
    \$\begingroup\$ @LinMa I don't have a lot of experience with Tries in Python, but I've seen a few implementations and all of them had a similar traverse() method. The search() method though is also needed to get us to the endpoint of a word from which we then traverse through different routes generating the resulting words. Personally, I am okay with this implementation, but I'd be happy to see if someone will come up with a better version. Not completely sure, but do you think you can post a separate question about this specific search+traverse set of methods? Thanks! \$\endgroup\$ – alecxe Feb 17 '17 at 2:40

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.