Home > Mobile >  Python Trie implementation - Inserting elements into a trie recursively
Python Trie implementation - Inserting elements into a trie recursively

Time:01-25

I have implemented a trie tree using a list in Python which stores words.

I am trying to write a function that insert words into a trie recursively.

Here my trie class and the function insert() :

class trie :
    def __init__(self, char):
        self.char = char
        self.children = []
    def __repr__ (self):
        return "%s %s" %(self.char, self.children)
    def __str__ (self):
        return "[%s %s]" %(self.char, self.children)

def rest(lst) : 
    return lst.children[-1]

def insert(root, word) :
    if len(word) == 0 :
        return "FIN"
    elif root.char == word[0] :
        return insert(rest(root), word[1:])
    elif root.char != word[0] :
        root.children.append(trie(word[0]))
    return insert(rest(root), word[1:])

The problem is that the insert() function does not insert the words into the correct childrens. For example :

t = trie("#") # root
insert(t, "cat")
insert(t, "#card")
insert(t, "#dog")
print(t)

The function insert() return the tree [c [a [t [r [d []]]], d [o [g []]]]], but the tree should be [c [a [t []], [r [d []]]], [d [o [g []]]]]. More specially, the character "r" and "d" should be in the children of "a".

CodePudding user response:

In the case that you find root.char == word[0], you have to pick the correct next child, not the last child. So the rest function only works when you're creating a new trie, but you need to check the char attributes in all children in the other case.

Below is an example. I wrote it iteratively since I think it's a bit simpler:

def insert(root, word):
    cur_root = root
    for ch in word:
        # Check all children to find the one which has the right char attr
        for child in cur_root.children:
            if child.char == ch:
                cur_root = child
                break
        else:
            # Not broken out of the loop, so no child was found with same char
            cur_root.children.append(trie(ch))
            cur_root = rest(cur_root)

CodePudding user response:

Your implementation has some issues:

  • It always adds new nodes below the last child of the grandparent
  • A Trie should give fast access to the child that corresponds to a given character. So it should be dictionary instead of a list. Or, if the set of characters is limited (like from "a" to "z") you could map a letter (using its ASCII code?) to a list index. But you should avoid the need to iterate a list.
  • A Trie node should not need a char attribute. You really need that character one step earlier, at the time you select that node. That is what a dictionary does: you first select by character, and then you get the corresponding node.

Here is an implementation just using dict:

def insert(root, word):
    if not word:
        root["FIN"] = True
    else:
        ch, *word = word
        if ch not in root:
            root[ch] = {} # Create child
        insert(root[ch], word)

def has(root, word):
    if not word:
        return "FIN" in root
    ch, *word = word
    return ch in root and has(root[ch], word)

t = {} # root
insert(t, "cat")
insert(t, "card")
insert(t, "dog")
print(t)  # {'c': {'a': {'t': {'FIN': True}, 'r': {'d': {'FIN': True}}}}, 'd': {'o': {'g': {'FIN': True}}}}

print(has(t, "ca"))  # False
print(has(t, "car"))  # False
print(has(t, "cat"))  # True
  • Related