Python Arena
Start learning
Step 6 of 6

Challenge: tries

A search box suggests words as you type: after "ca", it offers "cab", "cake", "car" and "card". With the words in a list, every question scans all of them with startswith: O(n·m) for n words and a prefix of m letters, again for each letter typed. A trie (from retrieval, said "try") stores the words letter by letter in a tree, so that words with the same beginning share it. Each node stands for a prefix. Its children are a dict from a letter to the next node, and a flag marks the nodes where a whole word ends. Here are "cab", "cake", "car", "card" and "dog", with a * where a word ends:

(root)
 +-- c
 |   +-- a
 |       +-- b *
 |       +-- k
 |       |   +-- e *
 |       +-- r *
 |           +-- d *
 +-- d
     +-- o
         +-- g *
class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_word = False


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

    def insert(self, word):
        node = self.root
        for letter in word:
            if letter not in node.children:
                node.children[letter] = TrieNode()
            node = node.children[letter]
        node.is_word = True

    def find(self, prefix):
        node = self.root
        for letter in prefix:
            if letter not in node.children:
                return None
            node = node.children[letter]
        return node

    def contains(self, word):
        node = self.find(word)
        return node is not None and node.is_word

    def complete(self, prefix):
        words = []
        node = self.find(prefix)
        if node is not None:
            self.collect(node, prefix, words)
        return words

    def collect(self, node, text, words):
        if node.is_word:
            words.append(text)
        for letter in sorted(node.children):
            self.collect(node.children[letter], text + letter, words)


trie = Trie()
for word in ["car", "dog", "card", "cake", "cab"]:
    trie.insert(word)
print(trie.contains("car"), trie.contains("ca"), trie.contains("cart"))
print(trie.complete("ca"))
print(trie.complete("do"), trie.complete("x"))
Output
True False False
['cab', 'cake', 'car', 'card']
['dog'] []

insert walks down the word, making the nodes that are missing, and marks the last one. find walks down a prefix one letter at a time and returns its node, or None when the path breaks off. "ca" is in the trie as a prefix, but no word ends there, which is why contains checks the flag.

Autocomplete in sorted order

Every word that starts with a prefix lies below the prefix's node. So complete finds that node, and collect visits everything under it in pre-order: a node's own word first, then its children. Two things make the words come out sorted, with no sorting of the words at the end. A word comes before the longer words that start with it ("car" before "card"), and the children are visited in alphabetical order.

That sorted matters. A dict keeps its keys in the order they were added, so a collect that loops over node.children as it is gives the words in an order that depends on the inserts:

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_word = False


def insert(root, word):
    node = root
    for letter in word:
        if letter not in node.children:
            node.children[letter] = TrieNode()
        node = node.children[letter]
    node.is_word = True


def collect(node, text, words):
    if node.is_word:
        words.append(text)
    for letter in node.children:
        collect(node.children[letter], text + letter, words)


root = TrieNode()
for word in ["car", "dog", "card", "cake", "cab"]:
    insert(root, word)
words = []
collect(root, "", words)
print(words)
Output
['car', 'card', 'cake', 'cab', 'dog']

Below "ca", the letters were added in the order r, k, b, because "car" came before "cake" and "cab", so "car" and "card" come out first. The same five words inserted in another order would come out in another order again. A test with words inserted in alphabetical order hides this mistake, which is why good tests mix them up.

The recursion in collect goes one level deeper per letter, so its depth is the length of the longest word: no danger of a RecursionError.

What it costs

insert, find and contains take one dict step per letter: O(m) for a word or a prefix of m letters, whether the trie holds ten words or a million. complete costs O(m) to find the node, then a visit to each node below it. Each visit builds its node's text with text + letter, which copies every letter so far. So with k nodes below the prefix and words of at most L letters, collect costs O(k·L): close to O(k) for short words, but in a trie of one word of 500 letters, complete("") copies over 125,000 letters. A version that keeps one list of letters, appends a letter on the way down, pops it on the way back up and calls "".join only where a word ends, costs O(k) plus the letters of the words it returns. The memory is one node per different prefix: at most the total number of letters in all the words, and less where words share their beginnings.

A trie can count, too. If every node keeps a count of the words that pass through it (insert adds 1 to the root and to each node on its way down), the number of words that start with a prefix is the count of the prefix's node: one walk of m letters instead of a scan of every word. The root is the node of the empty prefix "", which every word starts with, so its count is the number of words.

Your task

Challenge 1 of 3

A search engine shortens each word to the shortest stem it knows that the word starts with, so that "raining" and "rainbows" both become "rain". The stems are in a trie. Fill in the blanks of shortest_stem: it walks down the word one letter at a time, and at the first node where a stem ends, it returns the part of the word read so far, word[0] to word[i]. A word that starts with no stem stays as it is. The program should print:

['play', 'cat', 'saw', 'rain', 'in', 'the', 'rain']
ca rain

Challenge 2

Write count_with_prefix(words, prefixes) that returns a list with, for each prefix, the number of words that start with it. A word that is in the list twice counts twice, and every word starts with the empty prefix "". The hidden test has 60,000 words and 60,000 prefixes, so a scan of all the words for each prefix is too slow: build a trie whose nodes count the words that pass through them.

For example, this program:

words = ["car", "card", "care", "cat", "dog", "car"]
print(count_with_prefix(words, ["car", "ca", "d", "cart", "x"]))

prints:

[4, 5, 1, 0, 0]

Challenge 3

A word game lets players search with a dot for any one letter: "c.t" matches "cat", "cot" and "cut", but not "cart" (a dot is exactly one letter). Add a method matches(self, pattern) to the Trie class that returns the words of the trie that match pattern, in alphabetical order.

Walk down the trie along the pattern. A letter follows its one child, if there is one. A dot follows every child, in alphabetical order. When the whole pattern is used up, the node you're at is a match if a word ends there. A helper method that does this recursively, given the node, the pattern and the text so far, keeps it short.

For example, this program:

trie = Trie()
for word in ["cat", "cot", "cut", "cart", "dog", "dot", "cab"]:
    trie.insert(word)
print(trie.matches("c.t"))
print(trie.matches("do."), trie.matches("x.."))

prints:

['cat', 'cot', 'cut']
['dog', 'dot'] []
Previous: Heaps and heapq