Implement Trie (Prefix Tree)

Difficulty: Medium Source: NeetCode

Problem

A trie (pronounced as “try”) or prefix tree is a tree data structure used to efficiently store and retrieve keys in a dataset of strings. Implement the Trie class:

• Trie() Initializes the trie object.
• void insert(String word) Inserts the string word into the trie.
• boolean search(String word) Returns true if the string word is in the trie (i.e., was inserted before), and false otherwise.
• boolean startsWith(String prefix) Returns true if there is a previously inserted string that has the prefix prefix, and false otherwise.

Example 1: Input: ["Trie","insert","search","search","startsWith","insert","search"] with args [[],["apple"],["apple"],["app"],["app"],["app"],["app"]] Output: [null,null,true,false,true,null,true]

Constraints:

• 1 <= word.length, prefix.length <= 2000
• word and prefix consist only of lowercase English letters
• At most 3 * 10^4 calls total to insert, search, and startsWith

Prerequisites

Before attempting this problem, you should be comfortable with:

• Hash Maps — used as the children mapping at each trie node
• Trees — a trie is a tree where each node represents a character
• String traversal — processing a string character by character

1. Brute Force

Intuition

The simplest approach is to just store every inserted word in a Python set. Then search is an O(1) set lookup, and startsWith scans every stored word to check if any begins with the given prefix. This is easy to implement but startsWith can get slow as the word list grows.

Algorithm

1. Keep a set of inserted words.
2. insert(word): add word to the set.
3. search(word): return word in self.words.
4. startsWith(prefix): iterate through all stored words and return True if any starts with prefix.

Solution

class Trie:
    def __init__(self):
        self.words = 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)


trie = Trie()
trie.insert("apple")
print(trie.search("apple"))    # True
print(trie.search("app"))      # False
print(trie.startsWith("app"))  # True
trie.insert("app")
print(trie.search("app"))      # True

Complexity

• Time: O(n) for insert/search, O(W * n) for startsWith where W = number of words
• Space: O(W * L) where L = average word length

2. Trie with TrieNode

Intuition

A real trie stores characters along edges of a tree. Each node holds a dictionary of its children (one entry per character) and a boolean flag is_end that marks whether a complete word ends at that node. To insert, walk the tree creating nodes as needed. To search or check a prefix, walk the tree following existing nodes — the only difference between search and startsWith is whether you also check is_end at the final node.

Algorithm

1. Define TrieNode with children = {} and is_end = False.
2. insert(word): start at root, for each char create a child node if missing, then move to it. Mark is_end = True on the last node.
3. search(word): walk the trie for each char; if a char is missing return False. After the last char, return node.is_end.
4. startsWith(prefix): same walk but return True regardless of is_end if all chars were found.

Solution

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


class Trie:
    def __init__(self):
        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 search(self, word: str) -> bool:
        node = self.root
        for ch in word:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return node.is_end

    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return False
            node = node.children[ch]
        return True


trie = Trie()
trie.insert("apple")
print(trie.search("apple"))    # True
print(trie.search("app"))      # False
print(trie.startsWith("app"))  # True
trie.insert("app")
print(trie.search("app"))      # True
print(trie.startsWith("xyz"))  # False

Complexity

• Time: O(L) per operation where L = length of word/prefix
• Space: O(N * L) where N = number of inserted words

Common Pitfalls

Forgetting is_end on search. Walking to the last node of “app” succeeds even if only “apple” was inserted — you must check is_end to distinguish a prefix from a full stored word.

Mutating children between nodes. Each TrieNode must have its own children dict. If you write children = {} as a class variable instead of in __init__, all nodes share the same dict and everything breaks.

startsWith vs search confusion. startsWith returns True as long as the walk completes without a missing character — it does not care about is_end at all.