Implement Trie (Prefix Tree) — InterviewPrepKit

Solving tips

Build a tree of nodes each holding a children map (letter to node) plus an is_end boolean; words sharing a prefix share storage.
insert walks the path creating missing nodes and sets is_end on the last; search and startsWith both walk the path, differing only in the final is_end check.
Crucial distinction: search must also verify is_end, else 'app' wrongly matches after only 'apple' was inserted.
Every operation is O(L) in the word/prefix length, independent of how many words are stored; space is O(total characters inserted).

Problem

Build a data structure that stores a collection of lowercase strings and supports three operations:

Trie() — create an empty structure.
insert(word) — add word to the structure.
search(word) — return True if word was previously inserted as a complete word, False otherwise.
startsWith(prefix) — return True if any previously inserted word begins with prefix, False otherwise.

The distinction between search and startsWith matters: after inserting "apple", search("app") is False but startsWith("app") is True.

Examples

insert("apple"), then search("apple") → True — the exact word was inserted.
After the above, search("app") → False but startsWith("app") → True — "app" is only a prefix, not a stored word.
Then insert("app"), and now search("app") → True — the prefix became a full word too.

Constraints

1 <= len(word), len(prefix) <= 2000
All strings consist of lowercase English letters a–z only.
At most 3 * 10^4 calls total to insert, search, and startsWith.

Think about it first

Hint 1

Storing the words in a hash set makes search trivial — but how would you answer startsWith without scanning every stored word?

Hint 2

Words that share a prefix can share storage. Picture a tree where the root is the empty string and each edge is labeled with one letter; a word is the path of letters from the root.

Hint 3

Give each node a dict from letter to child node plus a boolean is_end. insert walks the path creating missing nodes and marks the last one; search and startsWith both walk the path, differing only in whether they check is_end at the end.

TL;DR

Trie of linked nodes: every operation is O(L) time in the length L of the word/prefix; space is O(total characters inserted).

Approach 1 — Brute force (hash set + linear prefix scan)

Keep a set for exact-word lookups and scan every stored word for prefix queries.

class Trie:
    def __init__(self) -> None:
        self.words: set[str] = set()

    def insert(self, word: str) -> None:
        self.words.add(word)

    def search(self, word: str) -> bool:
        return word in self.words

    def startsWith(self, prefix: str) -> bool:
        return any(w.startswith(prefix) for w in self.words)

insert and search are O(L), but startsWith is O(N · L) for N stored words. With up to 3 * 10^4 operations and words up to 2000 characters, a stream of inserts followed by startsWith calls approaches ~10^4 · 10^4 · 10^3 character comparisons in the worst case — far too slow, and it also wastes space by never sharing common prefixes.

Approach 2 — Trie with dict children

The insight: words that share a prefix should share storage. Arrange nodes in a tree where each edge is one letter; then a prefix query is just “does this path exist?”, answered in O(L) regardless of how many words are stored. A boolean flag on each node distinguishes “a stored word ends here” from “this is merely a waypoint”.

class TrieNode:
    __slots__ = ("children", "is_end")

    def __init__(self) -> None:
        self.children: dict[str, "TrieNode"] = {}
        self.is_end: bool = False


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

    def insert(self, word: str) -> None:
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
        node.is_end = True

    def _walk(self, s: str) -> "TrieNode | None":
        node = self.root
        for ch in s:
            node = node.children.get(ch)
            if node is None:
                return None
        return node

    def search(self, word: str) -> bool:
        node = self._walk(word)
        return node is not None and node.is_end

    def startsWith(self, prefix: str) -> bool:
        return self._walk(prefix) is not None

Walkthrough of the example. insert("apple") creates the chain root → a → p → p → l → e and sets is_end = True on the e node. search("app") walks root → a → p → p: the path exists, but that second-p node has is_end = False, so the answer is False. startsWith("app") walks the same three nodes, finds the path, and returns True without looking at is_end. After insert("app") — which creates no new nodes, it just re-walks a, p, p and flips is_end on the last one — search("app") now returns True.

Complexity: every operation touches one node per character, so O(L) time each. Space is O(total characters across all inserted words) in the worst case (no shared prefixes), and less when prefixes are shared.

Approach 3 — Trie with fixed 26-slot arrays

The insight: the alphabet is exactly a–z, so each node can hold a fixed list of 26 child slots. Indexing by ord(ch) - ord('a') replaces hashing, which lowers constant factors and makes per-node memory predictable — the classic competitive-programming layout.

class TrieNode:
    __slots__ = ("children", "is_end")

    def __init__(self) -> None:
        self.children: list["TrieNode | None"] = [None] * 26
        self.is_end: bool = False


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

    def _index(self, ch: str) -> int:
        return ord(ch) - ord("a")

    def insert(self, word: str) -> None:
        node = self.root
        for ch in word:
            i = self._index(ch)
            if node.children[i] is None:
                node.children[i] = TrieNode()
            node = node.children[i]
        node.is_end = True

    def _walk(self, s: str) -> "TrieNode | None":
        node = self.root
        for ch in s:
            node = node.children[self._index(ch)]
            if node is None:
                return None
        return node

    def search(self, word: str) -> bool:
        node = self._walk(word)
        return node is not None and node.is_end

    def startsWith(self, prefix: str) -> bool:
        return self._walk(prefix) is not None

Same O(L) time per operation. Space per node is a constant 26 slots whether or not they are used, so this trades memory (sparse tries waste slots) for speed (no hashing). Both trie variants are accepted; the dict version is more Pythonic.

Common pitfalls

Returning True from search whenever the path exists — you must also check is_end, or search("app") wrongly succeeds after only "apple" was inserted.
Forgetting that inserting a word that is a prefix of an existing word creates no new nodes; the only change is flipping is_end on an existing node.
Marking is_end on the node before the last character (off-by-one in the walk) — set it on the node reached after consuming the final letter.
Reaching for any(w.startswith(prefix) ...) over a set: correct but O(N · L) per prefix query, which is the exact cost the trie exists to remove.

Pattern takeaway

A trie converts “does any stored string have this prefix?” from a scan over all strings into a single root-to-depth-L walk. The reusable skeleton — children map plus is_end flag, with insert/walk helpers — is the base layer of nearly every trie problem; harder ones (wildcard search, Word Search II) only change how you walk it.