LeetCode 208: Implement Trie Prefix Tree

Problem Restatement

We need to implement a Trie, also called a prefix tree.

The Trie must support three operations:

The official problem asks us to implement the Trie class with insert, search, and startsWith. The Trie stores lowercase English words and supports exact word lookup and prefix lookup.

Input and Output

Class shape:

class Trie:

    def __init__(self):
        ...

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

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

    def startsWith(self, prefix: str) -> bool:
        ...

Examples

Example:

trie = Trie()
trie.insert("apple")
trie.search("apple")
trie.search("app")
trie.startsWith("app")
trie.insert("app")
trie.search("app")

Results:

None
True
False
True
None
True

Why?

After inserting "apple", the exact word "apple" exists.

But "app" is only a prefix, not a complete inserted word yet.

So:

trie.search("app") == False

However, some word starts with "app", namely "apple".

So:

trie.startsWith("app") == True

After inserting "app", it becomes a full word.

So:

trie.search("app") == True

Understanding a Trie

A Trie stores strings by sharing common prefixes.

If we insert:

app
apple
apt
bat

The Trie looks conceptually like:

root
├── a
│   └── p
│       ├── p  end of "app"
│       │   └── l
│       │       └── e  end of "apple"
│       └── t  end of "apt"
└── b
    └── a
        └── t  end of "bat"

The words "app" and "apple" share the prefix "app".

That is the main reason a Trie is useful.

Key Design

Each Trie node needs two things:

The is_word flag is necessary because a prefix and a full word are different.

For example, after inserting only:

"apple"

the path:

a -> p -> p

exists, but "app" has not been inserted as a word.

So search("app") must return False.

Trie Node

We can define a node like this:

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

children maps each character to the next node.

Example:

node.children["a"]

means the child reached by character "a".

Algorithm: Insert

To insert a word:

1. Start at the root.
2. For each character in the word:
   • If the character does not exist as a child, create a new node.
   • Move to that child.
3. After the last character, mark the current node as a complete word.

Example:

insert("apple")

We create or follow this path:

a -> p -> p -> l -> e

Then mark the e node as is_word = True.

Algorithm: Search

To search for an exact word:

1. Start at the root.
2. For each character:
   • If the character is missing, return False.
   • Move to that child.
3. After consuming all characters, return whether the current node is marked as a complete word.

The final check matters.

search("app") should return False after inserting only "apple".

Algorithm: startsWith

To check a prefix:

1. Start at the root.
2. For each character:
   • If the character is missing, return False.
   • Move to that child.
3. If all prefix characters are found, return True.

Unlike search, this does not require is_word = True.

Shared Helper

Both search and startsWith need to walk down the Trie.

We can write a helper:

def find_node(self, text: str):
    node = self.root

    for ch in text:
        if ch not in node.children:
            return None

        node = node.children[ch]

    return node

This returns the node at the end of the path, or None if the path does not exist.

Correctness

For insert, the algorithm creates exactly the path described by the word. Each character moves one level deeper in the Trie. After the final character, the last node is marked with is_word = True. Therefore the Trie records that the full word exists.

For search, the algorithm follows the path for the given word. If any character is missing, no inserted word can match that exact string, so returning False is correct. If the path exists, the algorithm returns True only when the final node is marked as a complete word. This correctly distinguishes a full inserted word from a prefix.

For startsWith, the algorithm follows the path for the prefix. If any character is missing, no inserted word has that prefix. If the whole prefix path exists, at least one inserted word begins with that prefix, so returning True is correct.

Thus all three operations match the required Trie behavior.

Complexity

Let m be the length of the input word or prefix.

insert may create one new node per character.

search and startsWith only traverse existing nodes.

Implementation

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_word = 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_word = True

    def search(self, word: str) -> bool:
        node = self.find_node(word)

        return node is not None and node.is_word

    def startsWith(self, prefix: str) -> bool:
        node = self.find_node(prefix)

        return node is not None

    def find_node(self, text: str):
        node = self.root

        for ch in text:
            if ch not in node.children:
                return None

            node = node.children[ch]

        return node

Code Explanation

The node stores outgoing edges and the word-ending marker:

self.children = {}
self.is_word = False

The Trie starts with an empty root:

self.root = TrieNode()

During insertion, we walk one character at a time:

for ch in word:

If the character path does not exist, create it:

if ch not in node.children:
    node.children[ch] = TrieNode()

Then move down:

node = node.children[ch]

After all characters are inserted, mark the final node:

node.is_word = True

Exact search uses the helper and checks the word marker:

return node is not None and node.is_word

Prefix search only checks whether the path exists:

return node is not None

Compact Implementation With defaultdict

We can also use defaultdict to reduce insertion code, but the explicit version is clearer for learning.

from collections import defaultdict

class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.is_word = False

class Trie:

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

    def insert(self, word: str) -> None:
        node = self.root

        for ch in word:
            node = node.children[ch]

        node.is_word = True

    def search(self, word: str) -> bool:
        node = self._walk(word)

        return node is not None and node.is_word

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

    def _walk(self, text: str):
        node = self.root

        for ch in text:
            if ch not in node.children:
                return None

            node = node.children[ch]

        return node

The explicit dictionary version avoids accidentally creating nodes during lookup.

Testing

def run_tests():
    trie = Trie()

    trie.insert("apple")
    assert trie.search("apple") is True
    assert trie.search("app") is False
    assert trie.startsWith("app") is True

    trie.insert("app")
    assert trie.search("app") is True

    assert trie.search("appl") is False
    assert trie.startsWith("appl") is True
    assert trie.startsWith("banana") is False

    trie.insert("bat")
    assert trie.search("bat") is True
    assert trie.startsWith("ba") is True
    assert trie.search("ba") is False

    print("all tests passed")

run_tests()