Implementing a Trie to support autocomplete in Python

问题

I'm trying to implement a data structure that supports autocomplete on a website. I've managed to implement an iterative version of a Trie. It supports the two primary methods of adding and searching in a Trie. However now I need to add a method that returns all the words that begin with the following prefix. Can someone help me with this.

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        curr = self.root
        for letter in word:
            node = curr.children.get(letter)
            if not node:
                node = TrieNode()
                curr.children[letter] = node
            curr = node
        curr.end = True

    def search(self, word):
        curr = self.root
        for letter in word:
            node = curr.children.get(letter)
            if not node:
                return False
            curr = node
        return curr.end

    def all_words_beginning_with_prefix(self, prefix):
        #I'm not sure how to go about this one.

回答1:

You could just implement a generator that iterates over the Trie according to prefix the same way as other methods do. Once you've found the node at the end of the prefix you can use recursive generator with yield from to iterate over the sub trie while keeping track of the prefix and yielding it when terminal node is found:

class TrieNode:
    def __init__(self):
        self.end = False
        self.children = {}

    def all_words(self, prefix):
        if self.end:
            yield prefix

        for letter, child in self.children.items():
            yield from child.all_words(prefix + letter)

class Trie:
    # existing methods here
    def all_words_beginning_with_prefix(self, prefix):
        cur = self.root
        for c in prefix:
            cur = cur.children.get(c)
            if cur is None:
                return  # No words with given prefix

        yield from cur.all_words(prefix)

trie = Trie()
trie.insert('foobar')
trie.insert('foo')
trie.insert('bar')
trie.insert('foob')
trie.insert('foof')

print(list(trie.all_words_beginning_with_prefix('foo')))

Output:

['foo', 'foob', 'foobar', 'foof']

回答2:

class Trie:

    def __init__(self):
        """
        Initialize your data structure here.
        """
        self.root = {}
        self.end = "#"

    def words_with_prefix(self, prefix: str):
        '''
        return all possible words with common prefix
        '''
        node = self.root
        for c in prefix:
            if c not in node:
                return []
            node = node[c]
        ans = []
        self._words_with_prefix_helper(node, prefix, ans)
        return ans

    def _words_with_prefix_helper(self, node, prefix, ans):
        for k in node:
            if k == self.end:
                ans.append(prefix)
                continue
            self._words_with_prefix_helper(node[k], prefix + k, ans)

complete implementation

来源：https://stackoverflow.com/questions/46038694/implementing-a-trie-to-support-autocomplete-in-python