Autocomplete

A search box that suggests words as you type is the job tries are best known for. All the words that start with what you typed sit in one branch of the trie, below one node. Find that node, then visit the branch.

Walk, then collect

Autocomplete has two parts:

1. Walk the prefix, exactly as in a prefix search. If a letter is missing, there are no suggestions.
2. Visit every node below the prefix node, and collect the word of every tagged node. Visiting each node of a tree once is called a traversal; going as deep as possible before going back up is a depth-first traversal (TREE).

The walk never looks at the to branch: it cannot hold a word that starts with te. The work is the length of the prefix plus the size of the branch below it.

The nodes do not store their words, so the traversal builds each word as it goes: the word of a child is the parent’s word plus the letter on the edge. Here is a version with an explicit stack instead of recursion, so a very long word cannot reach the recursion limit:

def complete(self, prefix, k=None):
    node = self._walk(prefix)
    if node is None:
        return []
    out = []
    stack = [(node, prefix)]
    while stack:
        node, word = stack.pop()
        if node.end:
            out.append(word)
            if len(out) == k:
                break
        # push in reverse order, so 'a' comes off the stack first
        for ch in sorted(node.children, reverse=True):
            stack.append((node.children[ch], word + ch))
    return out

Why the order is alphabetical

Two facts give alphabetical order. First, a node’s word comes before every word below it, because a word comes before any longer word that starts with it (tea before team). The code collects a node before pushing its children. Second, children are visited in alphabetical order of their letters, and every word in the a branch comes before every word in the b branch, because they first differ at that letter.

So the depth-first order is the sorted order, and asking for only the first k words lets the walk stop early (the “At most 2” variant). Sorting the children costs a little at each node; a node has at most 26 children, so it is a constant. A dict keeps keys in the order they were inserted, not in alphabetical order, which is why the code sorts them.

The best few, not all

A real search box shows the few most popular words, not the first few in alphabetical order. Say each word has a count, the number of times people searched for it, and the box shows the 3 with the highest counts. Visiting the whole branch is now unavoidable: the most popular word could be anywhere in it. After typing one letter, the branch holds about a twenty-sixth of all the words.

The fix is to do the work once, when words are inserted. Every node keeps its own short list: the best 3 words in its branch. When a word is inserted, it can only enter the lists of the nodes on its own path, so the insert walks its path and offers the word to each list.

def insert(self, word, count):
    node = self.root
    for ch in word:
        node = node.children.setdefault(ch, TrieNode())
        # keep the 3 best (count, word) pairs: highest count,
        # then alphabetical order for equal counts
        node.best.append((-count, word))
        node.best.sort()
        del node.best[3:]
    node.end = True

def suggest(self, prefix):
    node = self._walk(prefix)
    return [] if node is None else [w for _, w in node.best]

Storing -count makes the plain sort put the highest count first, and for equal counts the words come out in alphabetical order. Sorting a list of at most 4 pairs is a constant amount of work.

Now a suggestion costs O(L): walk the prefix and read the list. An insert costs O(L · k) for lists of length k, and the memory grows by k entries per node. This is a common trade: more work and memory when storing, so that the frequent question is cheap. Practice TRIE-P3 builds this search box.

A sorted list does this too

For plain alphabetical suggestions, a sorted list of words is a strong rival. All words that start with a prefix sit together in the sorted list, and binary search finds where they begin: bisect.bisect_left(words, prefix). Read forward from there while the words still start with the prefix:

import bisect

def complete(words, prefix, k):  # words is sorted
    i = bisect.bisect_left(words, prefix)
    out = []
    while i < len(words) and words[i].startswith(prefix) and len(out) < k:
        out.append(words[i])
        i += 1
    return out

This costs O(L log n + k), uses far less memory than a trie of Python objects, and the list is quick to build with one sort. It is a good choice when the words rarely change. Tries win when words are added and removed often (inserting into a sorted list shifts O(n) items), when the counts change, and when the trie stores extra facts at each node, such as the best-3 lists.