How to implement the startsWith function of a trie in python

Question

I have the following implementation of a trie in python.

I want to implement the startsWith function that will return a list of all words that start with the prefix.

Currently it return tue or false. Im not sure how to get a list of words in my implementation. I believe I need some sort of dfs, but not sure how to implement it

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

As you have logic in search and startsWith that is common, I would put that common logic in a separate method (I'll call it locate). Then on the TreeNode class you could define a recursive method that finds (yields) all words rooted at that node.

Code:

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