Insert, search and delete

Search follows one edge per letter, like insert, but never creates anything. Prefix search is the same walk with an easier finish. Delete is the only operation that needs care: it must remove the nodes that no longer lead to any word, and no others.

Search and prefix search

To search for a word, start at the root and, for each letter, follow the edge with that letter. Two things can happen:

• Some node has no edge for the next letter. Then no stored word even starts with the letters read so far, and the answer is no. Step through “Search tin”.
• Every letter is found. Then the answer depends on the end tag of the last node: tagged means the word is stored (“Search ten”); not tagged means it is only the start of longer words (“Search te”).

A prefix search (often called starts_with) asks whether any stored word starts with the given letters. The walk is the same; the answer is yes as soon as every letter is found, tag or no tag. This needs one rule to hold, which delete must keep:

Delete without leaving waste

To delete a word, walk to its node and remove its end tag. That is enough to make search say no. But it can break the rule above: if the node has no children, nothing is below it any more, and starts_with would still say yes for the deleted word.

So after removing the tag, walk back up the path. While the current node has no tag and no children, remove it from its parent. Stop at the first node that has a tag or another child, or at the root. Step through “Delete in, inn”: deleting in removes only its tag, because inn is below it; deleting inn then removes three nodes.

To walk back up, remember the path on the way down: a list of (parent, letter) pairs. Nodes do not need a link to their parent. A recursive delete also works, because the recursion keeps the path for you; a word is rarely long enough to reach Python’s recursion limit of about 1,000.

The code

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

    def _walk(self, s):
        # the node for prefix s, or None
        node = self.root
        for ch in s:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

    def search(self, word):
        node = self._walk(word)
        return node is not None and node.end

    def starts_with(self, prefix):
        return self._walk(prefix) is not None

    def delete(self, word):
        path = []           # (parent, letter) pairs
        node = self.root
        for ch in word:
            nxt = node.children.get(ch)
            if nxt is None:
                return False    # not stored
            path.append((node, ch))
            node = nxt
        if not node.end:
            return False
        node.end = False
        # remove nodes with no tag and no children
        while path and not node.end and not node.children:
            parent, ch = path.pop()
            del parent.children[ch]
            node = parent
        return True

Every operation costs O(L) for a word of length L: delete walks down once and back up at most once. dict.get(ch) returns None when the key is missing, instead of raising an error.

Counting words that pass through

A useful extra field is a count on every node: how many stored words have this node on their path. Insert adds 1 to every node it passes; delete takes 1 away. The count of a node is then the number of stored words that start with its prefix, so “how many words start with te?” costs O(L) instead of a walk over the whole branch.

Counts also make delete simpler: when a child’s count drops to 0, no word passes through it, so remove it and stop. And if the same word can be stored more than once, keep a number of copies at the end node instead of a True or False tag.

def insert(self, word):
    node = self.root
    node.count += 1
    for ch in word:
        node = node.children.setdefault(ch, TrieNode())
        node.count += 1
    node.copies += 1

setdefault(ch, TrieNode()) returns the child for ch, first storing the new node if there was none. It builds a new TrieNode every time, even when it is not needed; the if ch not in version avoids that and is a little faster. Practice TRIE-P2 asks for counts and copies.