Cannot write a function to retrieve all words in a trie

user_3832802 0
Viewed 106

I have a following Trie implementation:

class TrieNode:
    def __init__(self):
        self.nodes = defaultdict(TrieNode)
        self.is_fullpath = False


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

    def insert(self, word):
        curr = self.root
        for char in word:
            curr = curr.nodes[char]
        curr.is_fullpath = True

I'm trying to write a method to retrieve a list of all words in my trie.

t = Trie()
t.insert('a')
t.insert('ab')
print(t.paths())  # ---> ['a', 'ab']

My current implementation looks like this:

def paths(self, node=None):
    if node is None:
        node = self.root
    result = []
    for k, v in node.nodes.items():
        if not node.is_fullpath:
            for el in self.paths(v):
                result.append(str(k) + el)
        else:
            result.append('')
    return result

But it does not seem to return full list of words.

trie data-structures python
1 Answers

Here are the issues in your code:

  • It doesn't look further when is_fullpath is True. But you should also look deeper (for longer words) in that case.

  • It should not check node.is_fullpath but v.is_fullpath.

  • result.append('') is not correct. It should be result.append(str(k))

So your for loop body could look like this:

if v.is_fullpath:
    result.append(str(k))
for el in self.paths(v):
    result.append(str(k) + el)

I would however do it like this:

Define this recursive generator method on your TrieNode class:

def paths(self, prefix=""):
    if self.is_fullpath:
        yield prefix
    for chr, node in self.nodes.items():
        yield from node.paths(prefix + chr)

Note how this passes the collected characters on the path to the recursive call. If at any time the is_fullpath boolean is True, we yield that path. Always we continue the search recursively via child nodes.

The method on the Trie class is then quite simple:

def paths(self):
    return list(self.root.paths())
user_54598390
Related