DSA Implementing a Trie

The problem

A trie is a tree of letter nodes that stores words one letter per level. Implementing a trie means building one class that supports four operations on that tree.

The operations are insert, search, starts_with and count_prefix. Search answers whether a full word was inserted. Starts_with answers whether any inserted word opens with a prefix. Count_prefix answers how many inserted words open with a prefix.

Each node carries a counter for the words passing through it. That counter is what makes count_prefix a single walk instead of a scan. The operation script below is the contract the class must meet.

script = [
    ("insert", "cat"),
    ("insert", "car"),
    ("insert", "card"),
    ("insert", "dog"),
]
checks = [
    ("search", "cat", True),
    ("search", "ca", False),
    ("search", "card", True),
    ("search", "cards", False),
    ("starts_with", "ca", True),
    ("starts_with", "do", True),
    ("starts_with", "z", False),
    ("count_prefix", "ca", 3),
    ("count_prefix", "car", 2),
    ("count_prefix", "dog", 1),
]
for action, word in script:
    print(action, word)
for kind, text, expected in checks:
    print(kind, text, "->", expected)

Read the checks in groups. The search rows separate full words from prefixes and from words that were never inserted. The count rows read 3, 2 and 1 for three prefixes of shrinking reach.

Intuition: keep the words in a plain list

The first design stores each word in a list. Search compares the target against stored words one by one. A prefix count checks every stored word for the prefix.

The design is correct. Its cost is the trap. Every query pays for a scan of the whole list, and the list only grows.

words = []

def count_prefix_list(prefix):
    checks = 0
    total = 0
    for word in words:
        checks += 1
        if word.startswith(prefix):
            total += 1
    return total, checks

for word in ["cat", "car", "card", "dog"]:
    words.append(word)
total, checks = count_prefix_list("ca")
print("words starting with ca:", total)
print("words checked:", checks)

Four stored words cost four checks for one answer. Four thousand stored words would cost four thousand checks for the same kind of answer. The stored words share long openings, and the list design uses none of that sharing.

Brute force: the list design, counted

The brute force is the list design written as a class. Each query loops over the stored words. A counter records every word comparison the queries make.

class ListDictionary:
    def __init__(self):
        self.words = []
        self.checks = 0

    def insert(self, word):
        self.words.append(word)

    def search(self, word):
        for stored in self.words:
            self.checks += 1
            if stored == word:
                return True
        return False

    def starts_with(self, prefix):
        for stored in self.words:
            self.checks += 1
            if stored.startswith(prefix):
                return True
        return False

    def count_prefix(self, prefix):
        total = 0
        for stored in self.words:
            self.checks += 1
            if stored.startswith(prefix):
                total += 1
        return total

d = ListDictionary()
for word in ["cat", "car", "card", "dog"]:
    d.insert(word)
print("search card:", d.search("card"))
print("search ca:", d.search("ca"))
print("starts_with do:", d.starts_with("do"))
print("count_prefix ca:", d.count_prefix("ca"))
print("word checks so far:", d.checks)

Four small queries spend 15 word checks between them. A failed search is the most expensive query, because only a full scan proves the word is absent. The count query always pays the full scan, even when the answer is small.

The optimal solution: one node per letter, a counter per node

The state is a tree where each node holds a dict of child nodes keyed by letter, a word-end flag, and a counter of the words passing through. Every operation walks one node per letter, so its cost follows the word length and never the number of stored words.

Insert creates missing nodes on the path and raises the counter of every node it touches. Search walks the word and reads the end flag on the final node. Count_prefix walks the prefix and reads the counter on the final node.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False
        self.prefix_count = 0


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

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

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

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

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

    def count_prefix(self, prefix):
        node = self.walk(prefix)
        if node is None:
            return 0
        return node.prefix_count

trie = Trie()
for word in ["cat", "car", "card", "dog"]:
    trie.insert(word)
checks = [
    ("search", "cat", True),
    ("search", "ca", False),
    ("search", "card", True),
    ("search", "cards", False),
    ("starts_with", "ca", True),
    ("starts_with", "do", True),
    ("starts_with", "z", False),
    ("count_prefix", "ca", 3),
    ("count_prefix", "car", 2),
    ("count_prefix", "dog", 1),
]
passed = 0
for kind, text, expected in checks:
    if kind == "search":
        got = trie.search(text)
    elif kind == "starts_with":
        got = trie.starts_with(text)
    else:
        got = trie.count_prefix(text)
    ok = got == expected
    passed += 1 if ok else 0
    print(kind, text, ": got", got, ", expected", expected, ", pass", ok)
print("Passed", passed, "of", len(checks))

All ten checks from the start of the article pass. The word ca fails search but passes starts_with, and that split is the end flag doing its job. The path for ca exists, but no inserted word ends there.

Walkthrough: inserting car, then card

The walk inserts car first and reads the nodes along its path. It then inserts card and reads the same path again. Shared nodes keep one identity, and their counters rise.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False
        self.prefix_count = 0


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

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

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

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

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

    def count_prefix(self, prefix):
        node = self.walk(prefix)
        if node is None:
            return 0
        return node.prefix_count

trie = Trie()
trie.insert("car")
print("after inserting car:")
for label in ["", "c", "ca", "car"]:
    node = trie.walk(label)
    print("path", repr(label), "- count", node.prefix_count, "- end", node.is_end)
trie.insert("card")
print("after inserting card:")
for label in ["", "c", "ca", "car", "card"]:
    node = trie.walk(label)
    print("path", repr(label), "- count", node.prefix_count, "- end", node.is_end)

The node for car ends both readings with its flag set. Its counter rises from 1 to 2, because card also passes through it. The new node for d starts at 1, since only card reaches it.

Complexity

Each operation walks one node per letter, so insert and search take time proportional to the word length, written O(L) for a word of L letters. The number of stored words does not appear in that cost.

The tree stores one node per distinct letter position across all words, so its space grows with the total distinct prefixes. Shared openings are stored once.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False
        self.prefix_count = 0


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

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

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

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

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

    def count_prefix(self, prefix):
        node = self.walk(prefix)
        if node is None:
            return 0
        return node.prefix_count

def count_nodes(node):
    total = 1
    for child in node.children.values():
        total += count_nodes(child)
    return total

trie = Trie()
letters = 0
for word in ["cat", "car", "card", "dog"]:
    trie.insert(word)
    letters += len(word)
print("letters inserted:", letters)
print("nodes stored:", count_nodes(trie.root))
for n in [4, 40, 400]:
    print("stored words", n, "- list scan checks up to", n, "- trie steps for a 3-letter word:", 3)

The four script words carry 13 letters but need only 9 nodes. The shared ca opening is stored once, and the saving grows as more words share openings. A list scan grows with the stored word count, while a trie walk stays at the word length.

Edge cases and pitfalls

An empty trie answers every query with the failure value. The walk finds no first letter, so search is False and the count is 0.

The empty prefix is a real case. Its walk ends at the root, so starts_with answers True. The root counter holds every inserted word, so count_prefix for the empty prefix is the word total.

Inserting the same word twice is the classic pitfall. The end flag changes nothing the second time, but every counter on the path rises again. The counts then describe inserts, not distinct words.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False
        self.prefix_count = 0


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

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

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

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

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

    def count_prefix(self, prefix):
        node = self.walk(prefix)
        if node is None:
            return 0
        return node.prefix_count

trie = Trie()
print("search on empty trie:", trie.search("cat"))
print("count on empty trie:", trie.count_prefix("ca"))
for word in ["cat", "car", "card", "dog"]:
    trie.insert(word)
print("search empty word:", trie.search(""))
print("starts_with empty prefix:", trie.starts_with(""))
print("count empty prefix:", trie.count_prefix(""))
print("search cards:", trie.search("cards"))
print("count cards:", trie.count_prefix("cards"))
trie.insert("car")
print("count car after a duplicate insert:", trie.count_prefix("car"))

The duplicate insert lifts the car count from 2 to 3, even though the stored word set did not change. A design that needs distinct-word counts must reject a duplicate before walking, or check the end flag first.