Tries: prefix trees for strings

A trie, or prefix tree, stores a set of strings one character at a time. Every node stands for a prefix: the path from the root down to it spells that prefix, each child extends it by one character, and a flag marks the nodes where a stored word ends. Words that begin alike share the nodes of their common beginning, so pan, pant and pants live on a single path. To look up a word of length L, or to ask whether any stored word starts with a text of length L, the trie walks L steps down from the root, however many words it holds. Put a counter on every node and the same walk also tells how many stored words start with the text. Edward Fredkin described the structure in 1960 under the name “trie memory”, from retrieval. This lesson builds a trie, counts and lists words with it, measures what its nodes cost in CPython and then compresses it into a radix tree.

Words that begin alike share nodes

Take six words: pants, pan, pen, pine, pin and pant. Between them they have 22 letters, and the trie that stores them has 10 nodes below the root. The node p belongs to all six words; pa belongs to pan, pant and pants; pi belongs to pin and pine.

Three facts about the drawing carry the rest of the lesson.

• A node is a prefix, not a letter. The node pan stands for the text “pan”. It stores no letters of its own: they are written on the edges that lead down to it.
• The flag, not the shape, says which prefixes are words. pan ends a word and still has a child; pi has a child and ends nothing. You cannot tell the two apart without the flag.
• There is exactly one node for each different prefix, plus the root for the empty prefix. So a trie never has more nodes than its words have letters, and it has far fewer when the words share their beginnings.

Insert, search and starts_with

All three operations walk down from the root, one character at a time.

1. insert(word) follows the child for each character and creates it when it is missing. After the last character it sets the flag, even when every node on the path existed already.
2. search(word) follows the same children but creates nothing. A missing child means the word is not stored; at the end of the word, the flag gives the answer.
3. starts_with(prefix) walks the same way and answers yes as soon as the prefix is used up, because every node lies on the path of at least one stored word.

"""A trie of six words: insert, search and starts_with, and the nodes each insert creates."""


class Node:
    __slots__ = ("children", "is_word")

    def __init__(self):
        self.children = {}    # next character -> child node
        self.is_word = False  # does a stored word end at this node?


class Trie:
    def __init__(self):
        self.root = Node()  # the node of the empty prefix
        self.nodes = 0      # nodes below the root

    def insert(self, word):
        """Walks the word's path, adding the nodes that are missing; returns how many it added."""
        node, added = self.root, 0
        for ch in word:
            child = node.children.get(ch)
            if child is None:
                child = node.children[ch] = Node()
                added += 1
            node = child
        node.is_word = True  # even when no node was added
        self.nodes += added
        return added

    def _node_of(self, text):
        """The node whose path spells text, or None when no stored word starts with text."""
        node = self.root
        for ch in text:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

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

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


trie = Trie()
words = ["pants", "pan", "pen", "pine", "pin", "pant"]
for w in words:
    print(f"insert {w:6} adds {trie.insert(w)} node(s)")
letters = sum(len(w) for w in words)
print(f"{len(words)} words, {letters} letters, {trie.nodes} nodes below the root")
print()
for text in ["pan", "pa", "pine", "pines", "po", ""]:
    print(f"{text!r:8} search: {trie.search(text)!s:5}  starts_with: {trie.starts_with(text)}")

insert pants  adds 5 node(s)
insert pan    adds 0 node(s)
insert pen    adds 2 node(s)
insert pine   adds 3 node(s)
insert pin    adds 0 node(s)
insert pant   adds 0 node(s)
6 words, 22 letters, 10 nodes below the root

'pan'    search: True   starts_with: True
'pa'     search: False  starts_with: True
'pine'   search: True   starts_with: True
'pines'  search: False  starts_with: False
'po'     search: False  starts_with: False
''       search: False  starts_with: True

// A trie of six words: insert, search and startsWith, and the nodes each insert creates.

class TrieNode {
  constructor() {
    this.children = new Map(); // next character -> child node
    this.isWord = false; // does a stored word end at this node?
  }
}

class Trie {
  constructor() {
    this.root = new TrieNode(); // the node of the empty prefix
    this.nodes = 0; // nodes below the root
  }

  /** Walks the word's path, adding the nodes that are missing; returns how many it added. */
  insert(word) {
    let node = this.root;
    let added = 0;
    for (const ch of word) {
      let child = node.children.get(ch);
      if (child === undefined) {
        child = new TrieNode();
        node.children.set(ch, child);
        added++;
      }
      node = child;
    }
    node.isWord = true; // even when no node was added
    this.nodes += added;
    return added;
  }

  /** The node whose path spells text, or null when no stored word starts with text. */
  nodeOf(text) {
    let node = this.root;
    for (const ch of text) {
      node = node.children.get(ch);
      if (node === undefined) return null;
    }
    return node;
  }

  search(word) {
    const node = this.nodeOf(word);
    return node !== null && node.isWord;
  }

  startsWith(prefix) {
    return this.nodeOf(prefix) !== null;
  }
}

const trie = new Trie();
const words = ['pants', 'pan', 'pen', 'pine', 'pin', 'pant'];
for (const w of words) console.log(`insert ${w.padEnd(6)} adds ${trie.insert(w)} node(s)`);
const letters = words.reduce((n, w) => n + w.length, 0);
console.log(`${words.length} words, ${letters} letters, ${trie.nodes} nodes below the root`);
console.log();
for (const text of ['pan', 'pa', 'pine', 'pines', 'po', '']) {
  const shown = `'${text}'`.padEnd(8);
  console.log(`${shown} search: ${String(trie.search(text)).padEnd(5)}  startsWith: ${trie.startsWith(text)}`);
}

insert pants  adds 5 node(s)
insert pan    adds 0 node(s)
insert pen    adds 2 node(s)
insert pine   adds 3 node(s)
insert pin    adds 0 node(s)
insert pant   adds 0 node(s)
6 words, 22 letters, 10 nodes below the root

'pan'    search: true   startsWith: true
'pa'     search: false  startsWith: true
'pine'   search: true   startsWith: true
'pines'  search: false  startsWith: false
'po'     search: false  startsWith: false
''       search: false  startsWith: true

pants created five nodes, and pan, which came next, created none: its whole path was there already, so insert only set the flag on the node pan. pine added the nodes pi, pin and pine, after which pin and pant again only set flags. search("pa") answers no although the node pa exists, because no word ends there, while starts_with("pa") answers yes. The empty text is a special case. starts_with("") stops at the root before taking a step and answers yes, which is right here, since every word starts with the empty text, but it would also say yes for an empty trie. The counters in the next section fix that.

Two promises explain why the walks give the right answers. First, for every non-empty text t, the trie has a node for t exactly when some stored word begins with t. Second, a node’s flag is set exactly when its text was inserted. An empty trie, which is just the root, keeps both. Each insert adds nodes only for prefixes of its own word and sets only that word’s flag, so both promises still hold afterwards. starts_with relies on the first promise and search on both.

Counting the words below a prefix

A search box often needs a number rather than a yes or a no: how many stored words start with the text typed so far. Walking to the prefix’s node and then counting every word below it costs as much as that subtree is large. It is cheaper to keep two counters on each node: passing, the number of stored words whose path runs through the node, and ends, the number that end at it. insert adds 1 to passing on every node of its path, the root included, and 1 to ends at the last node. count_prefix(p) then walks down to the node of p and returns its passing, which takes as many steps as p has characters, or returns 0 when the walk breaks off. A word inserted twice counts twice, which is what a log of searches needs. starts_with becomes count_prefix(prefix) > 0, which also gives the right answer, no, for an empty trie.

"""Prefix counts on 2,000 made-up IFSC-style codes: branches per bank, and per block of branch numbers.

Each code has 11 characters, as an IFSC does: a 4-letter bank code, the character 0, then a 6-character branch code.
The bank codes and branch numbers are invented for this example.
"""
import random


class Node:
    __slots__ = ("children", "passing", "ends")

    def __init__(self):
        self.children = {}  # next character -> child node
        self.passing = 0    # stored codes whose path runs through this node
        self.ends = 0       # stored codes that end at this node


class CountingTrie:
    def __init__(self):
        self.root = Node()
        self.nodes = 0

    def insert(self, code):
        node = self.root
        node.passing += 1
        for ch in code:
            child = node.children.get(ch)
            if child is None:
                child = node.children[ch] = Node()
                self.nodes += 1
            node = child
            node.passing += 1
        node.ends += 1

    def _node_of(self, text):
        node = self.root
        for ch in text:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

    def search(self, code):
        node = self._node_of(code)
        return node is not None and node.ends > 0

    def count_prefix(self, prefix):
        """How many stored codes start with prefix: one walk of len(prefix) steps."""
        node = self._node_of(prefix)
        return 0 if node is None else node.passing


rng = random.Random(60)
banks = {"ALFA": 900, "BETA": 600, "GAMA": 350, "DLTA": 150}
codes = []
for bank, branches in banks.items():
    for number in rng.sample(range(1_000_000), branches):
        codes.append(f"{bank}0{number:06d}")
rng.shuffle(codes)

trie = CountingTrie()
for code in codes:
    trie.insert(code)

print(f"{len(codes):,} codes, {sum(map(len, codes)):,} characters, {trie.nodes:,} trie nodes")
print("first codes:", ", ".join(codes[:3]))
print()
for prefix in ["ALFA", "ALFA0", "ALFA05", "BETA099", "GAMA", "DLTA0", "ZETA", ""]:
    brute = sum(code.startswith(prefix) for code in codes)
    print(f"count_prefix({prefix!r:11}) = {trie.count_prefix(prefix):5,}   a scan of every code agrees: {trie.count_prefix(prefix) == brute}")
print()
print(f"search({codes[0]!r}) = {trie.search(codes[0])}, search('ALFA0') = {trie.search('ALFA0')}")

2,000 codes, 22,000 characters, 7,849 trie nodes
first codes: BETA0675603, ALFA0685611, ALFA0628572

count_prefix('ALFA'     ) =   900   a scan of every code agrees: True
count_prefix('ALFA0'    ) =   900   a scan of every code agrees: True
count_prefix('ALFA05'   ) =   111   a scan of every code agrees: True
count_prefix('BETA099'  ) =     6   a scan of every code agrees: True
count_prefix('GAMA'     ) =   350   a scan of every code agrees: True
count_prefix('DLTA0'    ) =   150   a scan of every code agrees: True
count_prefix('ZETA'     ) =     0   a scan of every code agrees: True
count_prefix(''         ) = 2,000   a scan of every code agrees: True

search('BETA0675603') = True, search('ALFA0') = False

Listing words in order: autocomplete

Autocomplete needs the words themselves: the first few stored words that start with the typed text, in dictionary order. Walk to the prefix’s node, then go depth first below it, trying the children in alphabetical order and reporting a word whenever the walk reaches a flagged node. A node’s own word is reported before the words below it, and that is exactly dictionary order, which puts pan before pant. Stop after k words.

"""Autocomplete over 10,000 made-up place names: the first 5 names with a prefix, in dictionary order."""
import random


class Node:
    __slots__ = ("children", "is_word")

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


def insert(root, word):
    node = root
    for ch in word:
        child = node.children.get(ch)
        if child is None:
            child = node.children[ch] = Node()
        node = child
    node.is_word = True


def first_names(root, prefix, k, in_order=True):
    """Up to k stored names that start with prefix, and the number of nodes the walk visited."""
    node = root
    for ch in prefix:
        node = node.children.get(ch)
        if node is None:
            return [], 0
    found, visited = [], 0
    stack = [(node, prefix)]
    while stack and len(found) < k:
        node, text = stack.pop()
        visited += 1
        if node.is_word:  # a name comes before every longer name that starts with it
            found.append(text)
        keys = sorted(node.children) if in_order else list(node.children)
        for ch in reversed(keys):  # pushed last, popped first: the smallest letter is explored first
            stack.append((node.children[ch], text + ch))
    return found, visited


rng = random.Random(15)
starts = ["a", "bi", "chi", "du", "ga", "ha", "ja", "ka", "ko", "la", "ma", "na", "pa", "ra", "sha", "si", "ta", "ve"]
consonants = ["d", "dh", "k", "l", "m", "n", "r", "s", "th", "v"]
vowels = ["a", "e", "i", "o", "u"]
ends = ["bad", "ganj", "gaon", "garh", "halli", "khed", "kot", "nagar", "ner", "palli", "pet", "pur", "pura", "wadi"]
names = set()
while len(names) < 10_000:
    parts = [rng.choice(starts)]
    for _ in range(rng.randint(0, 2)):
        parts.append(rng.choice(consonants) + rng.choice(vowels))
    parts.append(rng.choice(ends))
    names.add("".join(parts))
names = sorted(names)  # a set's order changes from run to run; sorting first makes the shuffle repeatable
rng.shuffle(names)

root = Node()
for name in names:
    insert(root, name)

for prefix in ["sha", "kot", "mad", "ve", "zz"]:
    found, visited = first_names(root, prefix, 5)
    every = sorted(n for n in names if n.startswith(prefix))
    assert found == every[:5], prefix  # the first 5 of a sorted scan of all 10,000 names
    print(f"{prefix!r:6} {len(every):5} names start with it; first 5: {', '.join(found) or '(none)'}; nodes visited: {visited}")

everything, visited = first_names(root, "sha", len(names))
print(f"listing all {len(everything)} names that start with 'sha' visits {visited:,} nodes")
found, _ = first_names(root, "sha", 5, in_order=False)
print("without sorting the children, 'sha' gives:", ", ".join(found))

'sha'    562 names start with it; first 5: shabad, shadadhopura, shadagarh, shadakakhed, shadakot; nodes visited: 25
'kot'     58 names start with it; first 5: kothabad, kothadakhed, kothadhabad, kothadhuwadi, kothadupet; nodes visited: 26
'mad'    119 names start with it; first 5: madakipet, madalewadi, madalipalli, madamakhed, madamapalli; nodes visited: 30
've'     556 names start with it; first 5: vebad, vedabad, vedadhuner, vedagaon, vedakhed; nodes visited: 23
'zz'       0 names start with it; first 5: (none); nodes visited: 0
listing all 562 names that start with 'sha' visits 2,625 nodes
without sorting the children, 'sha' gives: shalimopura, shalimepet, shalipur, shalipura, shalipalli

Every prefix got its five suggestions within 23 to 30 node visits, whether 58 names start with it or 562, because the walk stops at the fifth name. Listing all 562 names that start with sha took 2,625 visits. Each visited node leads to at least one stored word, so the walk never wanders into a dead end: from one reported word to the next it visits at most as many nodes as the longest word has letters. Five suggestions therefore cost at most L + 5m visits for a prefix of L letters and words of up to m letters.

The last line of the run shows a trap. A Python dict keeps its keys in the order they were inserted, which here is the order in which the names arrived, so a walk that does not sort them returns five names that are not the first five. Sort the keys at each node you visit, as the program does, keep the children sorted, or use an array of children indexed by letter, whose slots are already in alphabetical order.

What a trie costs: steps and nodes

None of these costs depends on n, the number of stored words. Each step is a lookup of one character in a node’s dict, which takes constant time on average, as in any hash map.

Memory is where a trie is expensive, and it can be estimated in two steps: count the nodes, then multiply by the cost of a node. The number of nodes is at most the total number of letters, and for random words it can be predicted. At depth d, the n_d words with at least d letters choose among 26^d possible prefixes, so the expected number of different prefixes at that depth is

and the trie’s expected number of nodes is the sum of N_d over all depths. The program below makes that estimate for 100,000 random words of 3 to 10 letters, counts the nodes of the real trie, and measures its memory with tracemalloc for two kinds of children:

"""A trie of 100,000 random words: its nodes, estimated and counted, and its memory with two kinds of children.

The words are 3 to 10 random lowercase letters. tracemalloc counts the memory blocks Python hands out while each
structure is built, so the figures are the same on every run of the same 64-bit CPython.
"""
import math
import random
import tracemalloc

LETTERS = "abcdefghijklmnopqrstuvwxyz"
SLOT = {ch: i for i, ch in enumerate(LETTERS)}


class DictNode:
    __slots__ = ("children", "is_word")

    def __init__(self):
        self.children = {}  # only the letters that occur
        self.is_word = False


class ListNode:
    __slots__ = ("children", "is_word")

    def __init__(self):
        self.children = [None] * 26  # one slot per letter, used or not
        self.is_word = False


def build_dict_trie(words):
    root, nodes = DictNode(), 0
    for w in words:
        node = root
        for ch in w:
            child = node.children.get(ch)
            if child is None:
                child = node.children[ch] = DictNode()
                nodes += 1
            node = child
        node.is_word = True
    return root, nodes


def build_list_trie(words):
    root, nodes = ListNode(), 0
    for w in words:
        node = root
        for ch in w:
            i = SLOT[ch]
            child = node.children[i]
            if child is None:
                child = node.children[i] = ListNode()
                nodes += 1
            node = child
        node.is_word = True
    return root, nodes


def expected_nodes(lengths, alphabet=26):
    """Expected distinct prefixes: at depth d, n_d random words fall into alphabet**d possible prefixes."""
    total = 0.0
    for d in range(1, max(lengths) + 1):
        n_d = sum(1 for n in lengths if n >= d)
        cells = alphabet**d
        total += cells * -math.expm1(n_d * math.log1p(-1 / cells))  # cells × (1 − (1 − 1/cells) ** n_d)
    return total


def traced(make):
    """What make() returns, and the bytes of memory it still holds afterwards."""
    tracemalloc.start()
    result = make()
    size = tracemalloc.get_traced_memory()[0]
    tracemalloc.stop()
    return result, size


rng = random.Random(26)
words, words_bytes = traced(lambda: ["".join(rng.choices(LETTERS, k=rng.randint(3, 10))) for _ in range(100_000)])
letters = sum(map(len, words))
print(f"{len(words):,} words ({len(set(words)):,} different), {letters:,} letters")
print(f"estimated nodes: {expected_nodes([len(w) for w in words]):,.0f}")
(_, nodes), dict_bytes = traced(lambda: build_dict_trie(words))
print(f"counted nodes:   {nodes:,}, one for every {letters / nodes:.2f} letters")
print()
(_, _), list_bytes = traced(lambda: build_list_trie(words))
_, set_bytes = traced(lambda: set(words))
print("memory held, measured with tracemalloc:")
print(f"  trie, a dict of children per node   {dict_bytes / 1e6:6.1f} MB   {dict_bytes / nodes:4.0f} bytes a node")
print(f"  trie, a list of 26 slots per node   {list_bytes / 1e6:6.1f} MB   {list_bytes / nodes:4.0f} bytes a node")
print(f"  the words as strings in a list      {words_bytes / 1e6:6.1f} MB")
print(f"  a set of those strings              {set_bytes / 1e6:6.1f} MB more")

100,000 words (96,199 different), 649,618 letters
estimated nodes: 359,731
counted nodes:   359,595, one for every 1.81 letters

memory held, measured with tracemalloc:
  trie, a dict of children per node     74.1 MB    206 bytes a node
  trie, a list of 26 slots per node    112.2 MB    312 bytes a node
  the words as strings in a list         5.6 MB
  a set of those strings                 4.2 MB more

The estimate of 359,731 nodes was within 0.04 % of the 359,595 counted. Most of those nodes are deep. The first three levels are crowded, with all 26 one-letter prefixes, all 676 two-letter ones and nearly all of the 17,576 three-letter ones present, but from the fifth letter on almost every word has a path of its own, so most nodes sit in chains of single children.

Then the bytes. In CPython a node with a dict of children took 206 bytes on average, and a node with a list of 26 slots took 312, because the list reserves a reference for every letter, 208 bytes on their own, while most nodes use one slot or none. Both nodes use __slots__, which stores their two fields in the object itself instead of in a dictionary of attributes. Even so, the trie of dicts held 74.1 MB for words whose strings take 5.6 MB, and a set of those strings costs only 4.2 MB more. A trie spends a lot of memory to answer questions about prefixes.

Array or hash-map children

The choice of children follows the alphabet. An array of σ slots per node pays for σ references whether they are used or not, and finds a child with a single index. A hash map pays only for the children that exist, at the price of hashing the character at every step. Arrays fit small, fixed alphabets such as digits, the four letters of DNA or the 26 lowercase letters, in languages where indexing an array is one memory access, such as C++, Java or Go. Maps fit large or open alphabets, such as all of Unicode, where an array per node would be almost entirely empty. In Python the list buys little speed, because every step still runs code to turn the letter into a slot number, and it costs half as much memory again, so a dict per node is the better default. The trie basics lesson of the coding patterns track measures the same choice in C++, where the array is the faster of the two.

Compressed tries: radix trees

In the trie of random words, most nodes have exactly one child and end no word. A compressed trie, also called a radix tree, merges every such chain into a single edge labelled with all of its letters. A node is kept only when it is the root, ends a word or has two or more children. NIST’s Dictionary of Algorithms and Data Structures describes the same idea under the name Patricia tree, after Donald Morrison’s PATRICIA of 1968, which compressed tries over the bits of keys. The six words of the first diagram keep seven of their ten nodes:

"""Compressing a trie: every chain of one-child nodes that ends no word becomes one edge with a longer label.

First the six words of the lesson's diagram, printed as a tree; then the 100,000 random words of the memory example,
with the nodes and the memory before and after.
"""
import random
import tracemalloc


class Node:
    __slots__ = ("children", "is_word")

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


class RadixNode:
    __slots__ = ("label", "children", "is_word")

    def __init__(self, label, is_word):
        self.label = label    # the letters on the edge into this node
        self.children = {}    # first letter of a child's label -> child
        self.is_word = is_word


def build_trie(words):
    root, nodes = Node(), 0
    for w in words:
        node = root
        for ch in w:
            child = node.children.get(ch)
            if child is None:
                child = node.children[ch] = Node()
                nodes += 1
            node = child
        node.is_word = True
    return root, nodes


def compress(node, label=""):
    """The radix tree of a trie node: one-child chains that end no word are merged into the edge above them."""
    out = RadixNode(label, node.is_word)
    for ch, child in node.children.items():
        text = ch
        while len(child.children) == 1 and not child.is_word:
            (ch2, child), = child.children.items()
            text += ch2
        out.children[ch] = compress(child, text)
    return out


def count(node):
    return 1 + sum(count(c) for c in node.children.values())


def show(node, depth=0):
    for ch in sorted(node.children):
        child = node.children[ch]
        print("  " * depth + f"{child.label!r}" + ("  (word)" if child.is_word else ""))
        show(child, depth + 1)


small, small_nodes = build_trie(["pants", "pan", "pen", "pine", "pin", "pant"])
tree = compress(small)
print(f"6 words: {small_nodes} trie nodes, {count(tree) - 1} after compression")
show(tree)

rng = random.Random(26)
LETTERS = "abcdefghijklmnopqrstuvwxyz"
words = ["".join(rng.choices(LETTERS, k=rng.randint(3, 10))) for _ in range(100_000)]
tracemalloc.start()
big, big_nodes = build_trie(words)
trie_bytes = tracemalloc.get_traced_memory()[0]
tracemalloc.stop()
tracemalloc.start()
radix = compress(big)
radix_bytes = tracemalloc.get_traced_memory()[0]
tracemalloc.stop()
radix_nodes = count(radix) - 1
print()
print(f"100,000 random words ({len(set(words)):,} different):")
print(f"  trie:            {big_nodes:7,} nodes  {trie_bytes / 1e6:5.1f} MB")
print(f"  compressed trie: {radix_nodes:7,} nodes  {radix_bytes / 1e6:5.1f} MB")
print(f"  nodes kept: {radix_nodes / big_nodes:.1%}, memory kept: {radix_bytes / trie_bytes:.1%}")

6 words: 10 trie nodes, 7 after compression
'p'
  'an'  (word)
    't'  (word)
      's'  (word)
  'en'  (word)
  'in'  (word)
    'e'  (word)

100,000 random words (96,199 different):
  trie:            359,595 nodes   74.1 MB
  compressed trie: 110,623 nodes   20.5 MB
  nodes kept: 30.8%, memory kept: 27.7%

On the random words, compression kept 110,623 of the 359,595 nodes, 30.8 %, and 27.7 % of the memory. A compressed trie of w different words has fewer than 2w nodes below the root: every one of them ends a word or branches, and the branching nodes that end no word are fewer than the leaves, each of which ends a word. The price is paid in the code. Search compares whole edge labels instead of single characters, which is still O(L) character comparisons, and insert must split an edge when a new word leaves it part of the way along: storing pa would split the edge an into a and n, with a new node pa between them.

Trie, hash set or sorted array?

A trie wins when prefix questions arrive while words are still being added, when each typed key should move one step down instead of starting a new search, and when the walk has to run in step with something else, such as a pattern with wildcards, a grid of letters or the bits of numbers, which is what the next lesson of this module does. For a list that does not change, a sorted array answers the same prefix questions in a small fraction of the memory.

Practise: a trie that counts

class Trie:
    """A trie of words that counts, for any prefix, how many inserted words start with it."""

    def __init__(self):
        pass  # set up the root node here

    def insert(self, word):
        """Stores one more copy of word."""

    def search(self, word):
        """True when word was inserted at least once."""
        return False

    def starts_with(self, prefix):
        """True when at least one inserted word starts with prefix."""
        return False

    def count_prefix(self, prefix):
        """How many inserted words, counting repeats, start with prefix."""
        return 0

from prefix_trie import Trie


def six_words():
    t = Trie()
    for w in ["pants", "pan", "pen", "pine", "pin", "pant"]:
        t.insert(w)
    return t


def test_search_needs_a_word_end():
    """search is True only for inserted words, not for every path"""
    t = six_words()
    assert t.search("pan") is True
    assert t.search("pants") is True
    assert t.search("pine") is True
    assert t.search("pa") is False
    assert t.search("p") is False
    assert t.search("pines") is False
    assert t.search("x") is False


def test_shorter_word_inserted_after_a_longer_one():
    """a word whose whole path already exists is still stored"""
    t = Trie()
    t.insert("pants")
    t.insert("pan")
    assert t.search("pan") is True
    assert t.search("pant") is False


def test_starts_with():
    """starts_with follows the path and needs no word end"""
    t = six_words()
    assert t.starts_with("pa") is True
    assert t.starts_with("pin") is True
    assert t.starts_with("pants") is True
    assert t.starts_with("po") is False
    assert t.starts_with("pantsx") is False


def test_count_prefix():
    """counts every inserted word below the prefix"""
    t = six_words()
    assert t.count_prefix("p") == 6
    assert t.count_prefix("pan") == 3
    assert t.count_prefix("pant") == 2
    assert t.count_prefix("pants") == 1
    assert t.count_prefix("pi") == 2
    assert t.count_prefix("pe") == 1
    assert t.count_prefix("pines") == 0
    assert t.count_prefix("q") == 0
    assert t.count_prefix("") == 6


def test_repeated_words_count_each_copy():
    """inserting a word twice counts it twice"""
    t = Trie()
    for w in ["ALFA0000123", "ALFA0000123", "ALFA0000456", "BETA0000123"]:
        t.insert(w)
    assert t.count_prefix("ALFA") == 3
    assert t.count_prefix("ALFA0000123") == 2
    assert t.count_prefix("") == 4
    assert t.search("ALFA0000123") is True


def test_empty_trie():
    """an empty trie stores no word, so nothing starts with the empty prefix"""
    t = Trie()
    assert t.search("a") is False
    assert t.starts_with("") is False
    assert t.count_prefix("") == 0


def test_many_codes():
    """50,000 codes and 50,000 counts, each count one walk down"""
    t = Trie()
    for i in range(50_000):
        t.insert(f"B{i % 7}{i:06d}")
    assert t.count_prefix("B3") == 7_143
    assert t.count_prefix("B6") == 7_142
    assert t.count_prefix("B30001") == 15
    assert t.count_prefix("B5049") == 143
    assert t.search("B3000101") is True
    assert t.search("B3000102") is False
    assert sum(t.count_prefix(f"B{i % 7}") for i in range(50_000)) == 357_142_858

/** A trie of words that counts, for any prefix, how many inserted words start with it. */
export class Trie {
  constructor() {
    // set up the root node here
  }

  /** Stores one more copy of word. */
  insert(word) {}

  /** true when word was inserted at least once. */
  search(word) {
    return false;
  }

  /** true when at least one inserted word starts with prefix. */
  startsWith(prefix) {
    return false;
  }

  /** How many inserted words, counting repeats, start with prefix. */
  countPrefix(prefix) {
    return 0;
  }
}

import { test, assert } from 'mysmartcopilot:test';
import { Trie } from './prefix_trie.mjs';

function sixWords() {
  const t = new Trie();
  for (const w of ['pants', 'pan', 'pen', 'pine', 'pin', 'pant']) t.insert(w);
  return t;
}

test('search is true only for inserted words, not for every path', () => {
  const t = sixWords();
  assert.equal(t.search('pan'), true);
  assert.equal(t.search('pants'), true);
  assert.equal(t.search('pine'), true);
  assert.equal(t.search('pa'), false);
  assert.equal(t.search('p'), false);
  assert.equal(t.search('pines'), false);
  assert.equal(t.search('x'), false);
});

test('a word whose whole path already exists is still stored', () => {
  const t = new Trie();
  t.insert('pants');
  t.insert('pan');
  assert.equal(t.search('pan'), true);
  assert.equal(t.search('pant'), false);
});

test('startsWith follows the path and needs no word end', () => {
  const t = sixWords();
  assert.equal(t.startsWith('pa'), true);
  assert.equal(t.startsWith('pin'), true);
  assert.equal(t.startsWith('pants'), true);
  assert.equal(t.startsWith('po'), false);
  assert.equal(t.startsWith('pantsx'), false);
});

test('countPrefix counts every inserted word below the prefix', () => {
  const t = sixWords();
  assert.equal(t.countPrefix('p'), 6);
  assert.equal(t.countPrefix('pan'), 3);
  assert.equal(t.countPrefix('pant'), 2);
  assert.equal(t.countPrefix('pants'), 1);
  assert.equal(t.countPrefix('pi'), 2);
  assert.equal(t.countPrefix('pe'), 1);
  assert.equal(t.countPrefix('pines'), 0);
  assert.equal(t.countPrefix('q'), 0);
  assert.equal(t.countPrefix(''), 6);
});

test('inserting a word twice counts it twice', () => {
  const t = new Trie();
  for (const w of ['ALFA0000123', 'ALFA0000123', 'ALFA0000456', 'BETA0000123']) t.insert(w);
  assert.equal(t.countPrefix('ALFA'), 3);
  assert.equal(t.countPrefix('ALFA0000123'), 2);
  assert.equal(t.countPrefix(''), 4);
  assert.equal(t.search('ALFA0000123'), true);
});

test('an empty trie stores no word, so nothing starts with the empty prefix', () => {
  const t = new Trie();
  assert.equal(t.search('a'), false);
  assert.equal(t.startsWith(''), false);
  assert.equal(t.countPrefix(''), 0);
});

test('50,000 codes and 50,000 counts, each count one walk down', () => {
  const t = new Trie();
  for (let i = 0; i < 50000; i++) t.insert(`B${i % 7}${String(i).padStart(6, '0')}`);
  assert.equal(t.countPrefix('B3'), 7143);
  assert.equal(t.countPrefix('B6'), 7142);
  assert.equal(t.countPrefix('B30001'), 15);
  assert.equal(t.countPrefix('B5049'), 143);
  assert.equal(t.search('B3000101'), true);
  assert.equal(t.search('B3000102'), false);
  let total = 0;
  for (let i = 0; i < 50000; i++) total += t.countPrefix(`B${i % 7}`);
  assert.equal(total, 357142858);
});

Interview questions

How does a trie count the stored words that start with a prefix in O(L) time? It keeps a counter on every node: the number of stored words whose path passes through that node. Inserting a word adds 1 to the counter of each node on its path, the root included, so a node’s counter always equals the number of stored words that start with its prefix. To answer a query, walk down from the root along the prefix, one character at a time; after L steps the counter of the node reached is the answer, and a missing child means the answer is 0. Nothing below the node is visited, so the cost is L steps whether 3 words or 3 million start with the prefix. A second counter of words that end exactly at a node lets the same trie answer search and handle repeated words.

When should trie nodes keep an array of children instead of a hash map? When the alphabet is small and fixed and speed matters most, such as digits, the four DNA letters or lowercase English in C++, Java or Go: an array finds the child with one index. The price is σ slots in every node, used or not, which is why the measured CPython trie took 312 bytes a node with a list of 26 slots against 206 with a dict. A hash map stores only the children that exist, so it suits large or open alphabets such as Unicode, where arrays would be mostly empty. In Python the dict is the better default: the list still needs code to turn each letter into a slot number, so it gains little speed for its extra memory.

Key takeaways

• A trie has one node per distinct prefix of its words; the path from the root spells the prefix, and a flag marks the prefixes that are words.
• insert, search and starts_with each walk L steps for a text of L characters, independent of how many words are stored. Set the flag even when insert creates no node.
• A passing counter on every node answers “how many words start with this prefix?” in the same L steps, and makes starts_with("") right for an empty trie.
• A depth-first walk that tries children in alphabetical order lists words in dictionary order; a Python dict keeps insertion order, so sort its keys.
• Estimate memory as nodes times bytes per node. In the CPython run, 100,000 random words needed 359,595 nodes, 206 bytes each with dict children and 312 with 26-slot lists.
• A compressed trie keeps only the root, word ends and branching nodes: 30.8 % of the nodes for the random words.