python 如何实现trie的startsWith函数

Question

我在 python.

中有一个 trie 的以下实现

我想实现 startsWith 函数，它将 return 以前缀开头的所有单词的列表。

目前 return 星期二或错误。我不确定如何在我的实现中获取单词列表。我相信我需要某种 dfs，但不确定如何实现它

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

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

    def insert(self, word):
        current = self.root
        for char in word:
            if char not in cur.children:
                current.children[char] = TrieNode()
            current = current.children[char]
        current.end_of_word = True
            
    
    def search(self, word):
        current = self.root
        for char in word:
            if char not in current.children:
                return False
            current = current.children[char]
        return current.end_of_word
        
    def startsWith(self, prefix):
        current = self.root
        for char in prefix:
            if char not in current.children:
                return False
            current = current.children[char]
        return True
    
    def remove(self, word):
        current = self.root
        if self.search(word):
            for char in word:
                current = current.children[char]
            current.end_of_word = False
        


# Your Trie object will be instantiated and called as such:
obj = Trie()
obj.insert('apple')
obj.insert('applegate')
obj.insert('apply')
obj.insert('application')

print(obj.search('applegate'))
obj.remove('applegate')
print(obj.search('applegate'))

print(obj.startsWith('app')) -> should return a list of all words that start with 'app'

Answer 1

你需要这样的东西：

def all_children_iterator(self, prefix: str) -> Iterator[str]:
    # returns all children of the current node with prefix prepended
    if self.end_of_word:
         yield prefix
    for char, child in self.children.items():
         yield from child.all_children_iterator(prefix + char)

那么你 return list(current.all_children_iterator(''))

而不是 return True

通过使用列表来保存字母并明智地附加和弹出它们，可以使这段代码的效率稍微高一些。但这使得基本思想有点难以理解。

Answer 2

因为您在 search 和 startsWith 中有共同的逻辑，所以我会将该共同逻辑放在一个单独的方法中（我将其称为 locate）。然后在 TreeNode class 上你可以定义一个递归方法来查找（产生）所有以该节点为根的单词。

代码：

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

    def allWords(self):
        if self.end_of_word:
            yield ""
        for char, child in self.children.items():
            for word in child.allWords():
                yield char + word

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

    def insert(self, word):
        current = self.root
        for char in word:
            if char not in current.children:
                current.children[char] = TrieNode()
            current = current.children[char]
        current.end_of_word = True
            
    def locate(self, word):
        current = self.root
        for char in word:
            if char not in current.children:
                return
            current = current.children[char]
        return current

    def search(self, word):
        current = self.locate(word)
        return current and current.end_of_word
        
    def startsWith(self, prefix):
        current = self.locate(prefix)
        if not current:
            return []
        return [prefix + word for word in current.allWords()]
    
    def remove(self, word):
        current = self.root
        if self.search(word):
            for char in word:
                current = current.children[char]
            current.end_of_word = False

python 如何实现trie的startsWith函数

How to implement the startsWith function of a trie in python

python

depth-first-search

trie

startswith