在 Python 中实现一个 Trie 以支持自动完成

Question

我正在尝试在网站上实现支持自动完成的数据结构。 我已经设法实现了 Trie 的迭代版本。它支持在 Trie 中添加和搜索的两种主要方法。 但是现在我需要添加一个方法来返回所有以以下前缀开头的单词。有人可以帮我解决这个问题吗？

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.

Answer 1

您可以只实现一个生成器，该生成器根据前缀对 Trie 进行迭代，就像其他方法一样。找到前缀末尾的节点后，您可以使用带有yield from 的递归生成器来迭代子树，同时跟踪前缀并在找到终端节点时生成它：

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')))

输出：

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

Answer 2

from collections import defaultdict


class TrieNode:
    def __init__(self):
        self.node = defaultdict(TrieNode)
        self.is_word = False


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

    def insert(self, words):
        curr = self.root

        for char in words:
            curr = curr.node[char]
        curr.is_word = True

    def search(self, word):
        curr = self.root
        for char in word:
            if char not in curr.node:
                return False
            curr = curr.node[char]
        return curr.is_word

    def dfs(self, node, word, word_list):
        if node.is_word == True:
            word_list.append(word)
        for a, n in node.node.items():
            self.dfs(n, word + a, word_list)

    def auto_complete(self, word_to_search, word_list):
        temp_word = ""
        curr = self.root
        for char in word_to_search:
            if char not in curr.node.keys():
                print("Invalid Input")
            else:
                temp_word += char
                curr = curr.node[char]
        self.dfs(curr, temp_word, word_list)
        print(word_list)

Answer 3

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