Problem: Given a trie (prefix tree), implement insertion and deletion of a word.

                    Root
        /             |            \
        A             I            T
        |          /    \      /   |    \
        #         N      S    A    H     O
                  |      |    |    |      \
                  #      #    V    E       W
                              |    | \     |
                              E    R  #    N
                              |    |       |
                              R    E       #
                              |    |
                              N    #
                              |
                              #

Input:

word: a string “[a-z]*”.

Output:

  • insert(word)
  • delete(word)

Definition of a TrieNode (TreeNode):

class TreeNode:
    def __init__(self, ch):
        self.val = ch
        self.children = [None]*27 # [0..25] are [a..z], [26] is for '#'.


Insertion (recursive):

1) If “”, reach the end of target word, and insert ‘#’ for representing that there is a word ending here.

2) If not “”,

  • 2.a) if child word[0] doesn’t exist. Create a new one.
  • 2.b) Insert word[1:] to the word[0]’s subtree.


Deletion (recursive):

1) What should be deleted? a node without any child (including ‘#’).

2) Where to delete? Node’s parent. A parent should get a return value to know whether or not it should delete that node.

Check whether a node should be deleted.

for child in root.children:
    if child != None:
        return False
return True

If target word is “”, then clean ‘#’. So that there is not a word anymore.

if not word:
    root.children[-1] = None

If target word is not empty, then decide whether we can delete child word[0].

ind = ord(word[0])-ord('a')
# the node exist, delete this node
if root.children[ind] and detele(root.children[ind], word[1:]):
    root.children[ind] = None

Deletion (iterative):

The iterative version is to use a stack, which stores the travel path. Then, it recursively deletes a pop-out node’s child, e.g., we delete ‘#’ for ‘E’ and ‘E’ for ‘H’.

If ‘H’ can be popped out from the stack, it means that ‘E’ can be deleted because it a null node with no children after deleting ‘#’. If ‘E’ has any other children than ‘#’, then we break the deleting loop. We shouldn’t delete ‘E’ because it is a part of another word.

The code is as follows:

def delete_iterative(root, word):
    """
    Iterative version: use stack.
        1) put the path into a stack.
        2) pop out and delete its child.
        3) check whether or not we can delete this node.
    """
    stack = [root]
    ptr = root
    for ch in word:
        ind = ord(ch)-ord('a')
        if ptr.children[ind]:
            stack.append(ptr.children[ind])
            ptr = ptr.children[ind]
        else:
            # not exist
            return

    n = len(stack) # = len(word)+1
    for i in xrange(n):
        ptr = stack.pop()

        if i == 0:
            ind = -1 # for '#'
        else:
            ind = ord(word[len(word)-i])-ord('a')

        ptr.children[ind] = None

        for child in ptr.children:
            if child:
                return


Complete solution (recursively):

class TreeNode:
    def __init__(self, ch):
        self.val = ch
        self.children = [None]*27

def insert(root, word):
    if not word:
        root.children[-1] = '#'
        return # easy to forget

    ind = ord(word[0])-ord('a')

    if not root.children[ind]:
        root.children[ind] = TreeNode(word[0])

    insert(root.children[ind], word[1:])

def detele(root, word):
    if not word:
        # We reach the end of target word:
        # e.g. "the" or "ther", to remove '#'.
        root.children[-1] = None
    else:
        # not reach the end.
        ind = ord(word[0])-ord('a')
        # the node exist, delete this node
        if root.children[ind] and detele(root.children[ind], word[1:]):
            root.children[ind] = None

    # answer whether or not this node can be removed.
    for child in root.children:
        if child != None:
            return False
    return True


def get_all(root, word, all_words):

    if root.children[-1] == '#':
        all_words.append(word)

    for i in range(26):
        if root.children[i] != None:
            new_word = word + root.children[i].val
            get_all(root.children[i], new_word, all_words)


root = TreeNode(None)

insert(root, "a")
insert(root, "in")
insert(root, "is")
insert(root, "tavern")
insert(root, "the")
insert(root, "there")
insert(root, "town")

words = []
get_all(root, "", words)
print words

words = []
detele(root, "ie")
detele(root, "ther")
detele(root, "there")
get_all(root, "", words)
print words

There are two questions on Leetcode relate to Trie: 208. Implement Trie (Prefix Tree) and 211. Add and Search Word - Data structure design.

In addition, Chapter 5.2 (page 730) in the Algorithms 4th introduces the Trie.