Trie in Python

Problem

The search box

The shop has 200,000 products. A customer types "co" in the search box, and suggestions must appear below it: "cola", "cord", "cocoa". With every new letter the suggestions change.

The first solution walks through all names and checks which ones start with what was typed. That is a linear search: 200,000 checks for every key pressed. A thousand customers typing at once make hundreds of millions of checks a second.

And most names do not even start with "c". Why look at them at all?

The idea

A tree of letters

Draw the names as paths. From the root, branches go out by the first letter, from each of them by the second, and so on. "Cola", "cocoa" and "cord" start with "co": those two letters are stored once, and the branches part only after them.

This tree is a trie (from retrieval, most often read as "try"). Every node keeps:

its children, by letter: a collection letter → node;
a mark that a name ends here. One name can end inside another: "cola" and "colander".

A prefix search has two parts:

Walking down the letters of the prefix, one step per letter. If a letter is missing, there are no such products.
Collecting: walking the branch under the last node. Every node marked "a name ends here" is a suggestion.

Walking down costs as much as the letters of the prefix, not as much as the products.

See it

Insert and prefix search

A dict keeps the children in insertion order. So collecting walks them with sorted, in alphabetical order.

class Node:
    def __init__(self):
        self.kids = {}
        self.end = False
        self.count = 0


class Trie:
    def __init__(self):
        self._root = Node()

    def insert(self, name, count):
        at = self._root
        for ch in name:
            if ch not in at.kids:
                at.kids[ch] = Node()
            at = at.kids[ch]
        at.end = True
        at.count = count

    def complete(self, prefix):
        at = self._root
        for ch in prefix:
            at = at.kids.get(ch)
            if at is None:
                return []
        found = []
        self._collect(at, prefix, found)
        return found

    def _collect(self, at, word, found):
        if at.end:
            found.append(word)
        for ch in sorted(at.kids):
            self._collect(at.kids[ch], word + ch, found)

Collecting is recursion, as in the lesson "Recursion": every node collects its names and asks its children. The base case is a node without children.

Four products and three searches:

trie = Trie()
trie.insert("cola", 120)
trie.insert("cocoa", 15)
trie.insert("cord", 40)
trie.insert("tea", 60)
print(f"co: [{', '.join(trie.complete('co'))}]")
print(f"cor: [{', '.join(trie.complete('cor'))}]")
print(f"m: [{', '.join(trie.complete('m'))}]")

co: [cocoa, cola, cord]
cor: [cord]
m: []

class Node:
    def __init__(self):
        self.kids = {}
        self.end = False
        self.count = 0


class Trie:
    def __init__(self):
        self._root = Node()

    def insert(self, name, count):
        at = self._root
        for ch in name:
            if ch not in at.kids:
                at.kids[ch] = Node()
            at = at.kids[ch]
        at.end = True
        at.count = count

    def complete(self, prefix):
        at = self._root
        for ch in prefix:
            at = at.kids.get(ch)
            if at is None:
                return []
        found = []
        self._collect(at, prefix, found)
        return found

    def _collect(self, at, word, found):
        if at.end:
            found.append(word)
        for ch in sorted(at.kids):
            self._collect(at.kids[ch], word + ch, found)

Step by step, in words

An empty trie: only the root. Every step down is one letter, so a path from the root spells a name.
insert("cola"): there are no letters yet, so every step makes a new node. The last one turns green: a name ends here.
"cocoa" starts like "cola": the first two letters already exist and are walked through. New nodes only where the names part.
"cord" shares the start too, and "tea" starts its own branch from the root. Four names, one shared start.
complete("co"): two steps down, one per letter of the prefix. That is the whole search, however many products there are.
Collecting: only the branch under the prefix is walked, children in alphabetical order. Every green node is a name and goes on the list.
Three completions: cocoa, cola, cord. "tea" was never looked at: its branch is not under the prefix.

"Tea" is never looked at: its branch is not under "co".

Try it

Predict the output

Three products. One name is the start of another:

trie = Trie()
trie.insert("cola", 1)
trie.insert("coat", 1)
trie.insert("colander", 1)
print(f"col: [{', '.join(trie.complete('col'))}]")

Answer

col: [coat, cola, colander]

col: [colander]

col: [cola, colander]

Try it

Fill in

What makes a node a suggestion while collecting?

if at.:
    found.append(word)

The idea

What it costs

Insert: one step per letter of the name: O(m), where m is its length.
Walking down a prefix: one step per letter of the prefix: O(m), however many products there are.
Collecting: as big as the branch under the prefix. If you show only the first 5 suggestions, you stop after the fifth.

Run the three ways. The catalogue is n names of 4 to 8 letters, in alphabetical order. The customer types 20 two-letter prefixes, and for each the first 5 names are needed. The count is the characters the search looks at: a comparison of a letter of the name with a letter of the prefix, or a step from a node to a child. The trie is built once, when the catalogue loads, and is not counted.

Scanning all names O(n)
Binary search in the sorted catalogue O(log n)
Trie O(1)

Algorithm	n = 10	n = 1,000	n = 10,000	n = 100,000
Scanning all names	223	11,778	93,071	1,303,482
Binary search in the sorted catalogue	125	542	676	808
Trie	46	403	342	266

The code runs in your browser, in JavaScript. The number of operations is the same in every language, because the algorithm is the same. The time depends on the machine, so it is only for comparison. “≈” is an estimate from the curve: we do not run an n that big, so the page does not freeze.

class Node:
    def __init__(self):
        self.kids = {}
        self.end = False
        self.count = 0


class Trie:
    def __init__(self):
        self._root = Node()

    def insert(self, name, count):
        at = self._root
        for ch in name:
            if ch not in at.kids:
                at.kids[ch] = Node()
            at = at.kids[ch]
        at.end = True
        at.count = count

    def complete(self, prefix):
        at = self._root
        for ch in prefix:
            at = at.kids.get(ch)
            if at is None:
                return []
        found = []
        self._collect(at, prefix, found)
        return found

    def _collect(self, at, word, found):
        if at.end:
            found.append(word)
        for ch in sorted(at.kids):
            self._collect(at.kids[ch], word + ch, found)

With 100,000 names the scan looks at 1,303,482 characters, binary search in the sorted catalogue at 808, and the trie at 266. The more names, the denser the tree and the sooner the five are collected: with 1000 names the trie looks at 403 characters, with 100,000 at fewer.

Binary search is here because the catalogue is sorted: then it finds the first name with the prefix in log₂ n comparisons (binary search). See below for more.

See it

Most bought first

Alphabetical order is not the most useful. If a customer types "co", they are more likely after cola than cocoa. So every node where a name ends also keeps how many times the product was bought (count). The suggestions are sorted by it:

def count_of(self, name):
    at = self._root
    for ch in name:
        at = at.kids[ch]
    return at.count

def suggest(self, prefix, k):
    names = self.complete(prefix)
    return sorted(names, key=self.count_of, reverse=True)[:k]

trie = Trie()
trie.insert("cola", 120)
trie.insert("cocoa", 15)
trie.insert("cord", 40)
trie.insert("tea", 60)
trie.insert("colander", 8)
trie.insert("coat", 25)
print(f"co: [{', '.join(trie.suggest('co', 3))}]")

co: [cola, cord, coat]

The sort is stable (see stability): with an equal number of purchases the alphabetical order stays.

More precisely, if you are curious

Here every search collects all names under the prefix and then sorts them. For "c" in a real catalogue those are tens of thousands. That is why big search engines keep, in every node, the ready first 5–10 suggestions for the branch below it. The search becomes just walking down, and the price is paid on insert: every node on the path of a new name checks whether it gets into its top ones.

Try it

Find the bug

A colleague wrote the collecting, but the suggestions are full of pieces of words:

def _collect(self, at, word, found):
    found.append(word)
    for ch in sorted(at.kids):
        self._collect(at.kids[ch], word + ch, found)

trie = Trie()
trie.insert("cola", 120)
trie.insert("cocoa", 15)
trie.insert("cord", 40)
trie.insert("tea", 60)
print(f"co: [{', '.join(trie.complete('co'))}]")

co: [co, coc, coco, cocoa, col, cola, cor, cord]

Which line has the bug?

Pick the line with the bug.

The idea

Letter case and look-alike characters

For the trie "C" and "c" are different letters: different children of the node. A customer who types "Cola" or " cola" with a space in front finds nothing. So both the name on insert and the text on search go through one and the same normalisation: no spaces at the ends, one space between words, lower case.

def normalize(text):
    return " ".join(text.split()).lower()

trie = Trie()
trie.insert("cola", 120)
trie.insert("cocoa", 15)
trie.insert("cord", 40)
trie.insert("tea", 60)
trie.insert("colander", 8)
trie.insert("coat", 25)
trie.insert("café", 30)
print(f"CoLa: [{', '.join(trie.complete(normalize('  CoLa ')))}]")
print(f"cafe: [{', '.join(trie.complete(normalize('cafe')))}]")

CoLa: [cola, colander]
cafe: []

The second search finds nothing, although the catalogue has "café". "e" and "é" are different characters, different children in the tree. The same happens with letters that look exactly alike: the Latin "a" and the Cyrillic letter that looks exactly like it, in a shop that sells in two alphabets. Normalisation does not fix that by itself: you need an explicit replacement (for example removing accents) if you want to forgive it.

More precisely, if you are curious

lower() is enough for most alphabets. For German and some other languages there is casefold(): it also turns "ß" into "ss". A string in Python is walked by characters (code points), so an emoji is one character too.

How it is in the other languages: Java and C# walk by 16-bit char, Python and for…of in JavaScript by characters. For Latin and Cyrillic there is no difference.

The idea

Memory, and when a sorted array is enough

Every node keeps its children. There are two ways:

A collection letter → node, as here: the node keeps only the children it has.
An array of slots, one per letter: 26 for the Latin alphabet, 30 for the Bulgarian one, 32 with space and hyphen. The child is found at once by the letter's number, but every node keeps all the slots, the empty ones too.

How empty are they? The trie in Bench for 100,000 names (631,833 letters) has 210,642 nodes, on average one child per node: the shared start is short, then each name goes on its own. An array of 30 slots per node keeps more than 6.3 million references, of which 210,641 are used. The collection is smaller, and the array is faster only in the top levels, where children are many.

If the names change rarely and you only need "all that start with …", a sorted array is often enough. Binary search finds the first name that is not before the prefix, and the matches sit next to each other after it:

from bisect import bisect_left

names = ["coat", "cocoa", "cola", "colander", "cord", "tea"]
prefix = "co"
at = bisect_left(names, prefix)
found = []
while at < len(names) and names[at].startswith(prefix):
    found.append(names[at])
    at += 1
print(f"{prefix}: [{', '.join(found)}]")

co: [coat, cocoa, cola, colander, cord]

bisect_left returns the first place where the prefix can be inserted without breaking the order: that is the first name that is not before it.

A trie wins when names are added and removed all the time, when you need counts and the most bought for every prefix, or when you search letter by letter while the customer types: every new letter is one step from the node you have already reached.

Exercise

How many products start with…

Solve the task in your editor.

The count under a prefix. Next to every suggestion the box shows how many products start with what was typed. Add to the node a counter of how many names pass through it, and a method count_prefix that walks down the prefix and returns the counter, without collecting names. Insert the whole catalogue from the card "Most bought first" and check the prefixes "co", "cor", "t" and "m":

co: 5
cor: 1
t: 1
m: 0

More tasks

Deleting. A product stops selling. Remove the name from the trie without breaking the names that pass through the same nodes.
Number of nodes. Count the nodes after inserting the catalogue and compare with the total number of letters in the names.
Letter by letter. Make a search that keeps the node it has reached and, for every new letter, takes only one step.

Test

Check yourself

1. The catalogue has 200,000 products. How many steps does the trie take to reach the node of the prefix "cola"?200,000: one per product4: one per letterAbout 18: log₂ of 200,000

2. Why does a node have a separate "a name ends here" mark?Because leaves have no childrenBecause one name can end inside the path of another: "cola" is the start of "colander"To make walking down faster

3. The product names change once a day, and you only need "all that start with …". What is enough?A trie, no matter whatScanning all namesA sorted array and a binary search for the first name with the prefix

4. The customer types "Cola", and the catalogue has "cola". What do you do?Nothing: the search ignores caseNormalise both: lower case and no extra spaces, both when inserting and when searchingStore every product twice

Trie

Cookies