← Back to DevBytes

Solving Implement Trie (Prefix Tree) in Python: Step-by-Step Guide

Introduction to the Trie Data Structure

A Trie, pronounced "try" (from the word "retrieval"), is a specialized tree-based data structure designed for efficient storage and retrieval of strings. Unlike binary search trees or hash tables, a Trie organizes data character by character, making it exceptionally powerful for prefix-based operations. Each node in a Trie represents a single character, and a path from the root to any node represents a prefix shared by all strings that pass through that path.

The "Implement Trie (Prefix Tree)" problem is a classic algorithmic challenge that appears frequently in coding interviews and real-world applications. It typically requires you to build a Trie class that supports three core operations: inserting a word, searching for a complete word, and checking whether any word in the Trie starts with a given prefix. Mastering this problem unlocks a deeper understanding of how autocomplete systems, spell checkers, and IP routing tables work under the hood.

Why Tries Matter

Before diving into implementation, it is important to understand why Tries are valuable compared to other data structures. When you store a list of words in a standard array or hash set, searching for a word takes O(n) or O(1) time respectively, but prefix-based queries become expensive. For example, finding all words that start with "app" in a hash set requires scanning every entry, resulting in O(n * m) time complexity where n is the number of words and m is the average word length.

Tries solve this problem elegantly. Insertion and search operations both run in O(m) time, where m is the length of the word being inserted or searched. This time complexity is independent of the number of words stored in the Trie. Additionally, Tries naturally group words by their shared prefixes, which means memory usage can be significantly lower than storing each word independently when many words share common beginnings.

Real-World Applications

Understanding the Trie Structure

A Trie consists of nodes, where each node contains two key pieces of information: a collection of child nodes (one for each possible character) and a flag indicating whether that node marks the end of a complete word. The root node is special because it represents an empty string and has no character associated with it. From the root, each level of the tree corresponds to one character position in the strings being stored.

For example, if you insert the words "cat", "car", and "card" into a Trie, the structure would share the "ca" prefix across all three words, then branch at the third character. The "car" and "card" words would further share the "car" prefix before branching at the fourth character. This sharing is what makes Tries memory-efficient for datasets with overlapping prefixes.

Step-by-Step Implementation in Python

Now let us build a complete Trie implementation. We will create two classes: a TrieNode class to represent individual nodes, and a Trie class that exposes the public API with insert, search, and startsWith methods.

Step 1: Defining the TrieNode Class

Each node needs a dictionary to store its children, mapping characters to child TrieNode objects. We also need a boolean flag to mark whether a node represents the end of a word. Using a dictionary for children is flexible and works well for any character set, though for lowercase English letters only, you could use an array of size 26 for slightly faster access.

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

This simple class gives us the building block for our Trie. Every node starts with an empty dictionary of children and the end-of-word flag set to False. Only nodes that represent the final character of an inserted word will have this flag set to True.

Step 2: Initializing the Trie Class

The Trie class itself starts with a root node. This root node does not represent any character; it simply serves as the entry point from which all words branch out.

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

Step 3: Implementing the Insert Method

To insert a word, we start at the root and traverse character by character. For each character, we check if a child node exists for that character. If it does not, we create a new TrieNode and add it to the current node's children. We then move to that child node and repeat for the next character. After processing all characters, we mark the final node as the end of a word.

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

Notice how this method handles overlapping prefixes automatically. If you insert "apple" and then "app", the second insertion reuses the existing nodes for "a", "p", and "p", simply marking the third node as an end of word. No duplicate nodes are created for shared prefixes.

Step 4: Implementing the Search Method

Searching for a word follows a similar traversal pattern. We walk through each character of the word, moving down the Trie. If at any point a character does not exist in the current node's children, the word is not in the Trie and we return False. If we successfully traverse all characters, we check the end-of-word flag. This flag check is critical because inserting "apple" does not mean "app" exists as a complete word, even though all its characters exist in the Trie.

def search(self, word: str) -> bool:
    current = self.root
    for char in word:
        if char not in current.children:
            return False
        current = current.children[char]
    return current.is_end_of_word

Step 5: Implementing the startsWith Method

The startsWith method checks whether any word in the Trie begins with a given prefix. The logic is nearly identical to search, except we do not check the end-of-word flag at the end. If we can traverse all characters of the prefix without missing any node, then at least one word with that prefix exists in the Trie.

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

Complete Implementation

Putting all the pieces together, here is the complete Trie implementation ready for use:

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


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

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

    def search(self, word: str) -> bool:
        current = self.root
        for char in word:
            if char not in current.children:
                return False
            current = current.children[char]
        return current.is_end_of_word

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

Testing the Implementation

To verify that the implementation works correctly, let us write a series of test cases that exercise all three operations. These tests cover inserting words, searching for words that exist and do not exist, and checking prefixes.

# Create a new Trie instance
trie = Trie()

# Insert words
trie.insert("apple")
trie.insert("app")
trie.insert("application")
trie.insert("banana")
trie.insert("band")

# Test search
print(trie.search("apple"))       # True
print(trie.search("app"))         # True
print(trie.search("appl"))        # False (prefix exists but not a complete word)
print(trie.search("application")) # True
print(trie.search("banana"))      # True
print(trie.search("bandana"))     # False
print(trie.search("cat"))         # False

# Test startsWith
print(trie.startsWith("app"))     # True
print(trie.startsWith("appl"))    # True
print(trie.startsWith("ban"))     # True
print(trie.startsWith("band"))    # True
print(trie.startsWith("can"))     # False
print(trie.startsWith("xyz"))     # False

All outputs should match the expected boolean values shown in the comments. This confirms that the Trie correctly distinguishes between complete words and prefixes, and that shared prefixes are handled properly.

Complexity Analysis

Understanding the time and space complexity of each operation is essential for knowing when to use a Trie. Let m represent the length of the word or prefix being processed and n represent the total number of characters across all inserted words.

Compare this to using a hash set for the same operations. A hash set provides O(1) average-case search for exact words, but prefix queries degrade to O(n * m) because you must scan every word. The Trie trades slightly higher memory overhead for dramatically faster prefix operations.

Best Practices and Optimization Techniques

Use Arrays for Fixed Character Sets

If you know your Trie will only store lowercase English letters, replacing the dictionary with a fixed-size array of 26 elements can improve performance. Array indexing is faster than dictionary lookups, and the memory overhead per node becomes predictable.

class TrieNode:
    def __init__(self):
        self.children = [None] * 26
        self.is_end_of_word = False

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

    def insert(self, word: str) -> None:
        current = self.root
        for char in word:
            index = ord(char) - ord('a')
            if current.children[index] is None:
                current.children[index] = TrieNode()
            current = current.children[index]
        current.is_end_of_word = True

    def search(self, word: str) -> bool:
        current = self.root
        for char in word:
            index = ord(char) - ord('a')
            if current.children[index] is None:
                return False
            current = current.children[index]
        return current.is_end_of_word

    def startsWith(self, prefix: str) -> bool:
        current = self.root
        for char in prefix:
            index = ord(char) - ord('a')
            if current.children[index] is None:
                return False
            current = current.children[index]
        return True

Consider Memory Optimization with Compressed Tries

Standard Tries can waste memory when words have long unique suffixes with no sharing. A compressed Trie, also known as a Radix Tree, stores strings at edges rather than single characters. This reduces the number of nodes significantly for sparse datasets. While more complex to implement, Radix Trees are worth considering for memory-constrained environments.

Add a Delete Method for Completeness

Many Trie implementations benefit from a delete operation. Deleting a word requires traversing to the end node, unsetting the end-of-word flag, and then removing nodes that are no longer part of any other word. You must be careful not to remove nodes that are still used by other words or that are prefixes of other words.

def delete(self, word: str) -> None:
    def _delete_helper(node, word, index):
        if index == len(word):
            if not node.is_end_of_word:
                return False
            node.is_end_of_word = False
            return len(node.children) == 0

        char = word[index]
        if char not in node.children:
            return False

        should_delete_child = _delete_helper(node.children[char], word, index + 1)

        if should_delete_child:
            del node.children[char]
            return len(node.children) == 0 and not node.is_end_of_word

        return False

    _delete_helper(self.root, word, 0)

This recursive helper returns True when a node can be safely deleted, which happens only when it has no children and is not marked as the end of a word. The parent then removes the child reference, and the process bubbles up toward the root.

Validate Input Early

In production code, validate that inputs contain only expected characters before processing. This prevents subtle bugs and makes debugging easier. For the array-based implementation, an unexpected character like a digit or uppercase letter would cause an index error or silently produce incorrect results.

Use Type Hints for Clarity

Adding type hints makes your code more readable and enables better IDE support. This is especially valuable in a data structure implementation where the types of nodes and return values matter.

from typing import Dict

class TrieNode:
    def __init__(self) -> None:
        self.children: Dict[str, 'TrieNode'] = {}
        self.is_end_of_word: bool = False

Common Pitfalls to Avoid

One frequent mistake is forgetting to check the is_end_of_word flag in the search method. Without this check, searching for "app" after inserting only "apple" would incorrectly return True because all the characters exist in the Trie. The flag is what distinguishes a complete word from a mere prefix.

Another common error is confusing the search and startsWith methods. Remember that search requires the exact word to exist, while startsWith only requires the prefix path to exist. These are fundamentally different queries with different correctness criteria.

Finally, be cautious with recursive implementations for very long words. Python has a default recursion limit of around 1000, so words longer than that could cause a stack overflow. The iterative approach shown in this tutorial avoids this issue entirely.

Conclusion

Implementing a Trie in Python is a rewarding exercise that deepens your understanding of tree-based data structures and prefix matching. The implementation itself is straightforward once you grasp the core concept of character-by-character traversal, yet the resulting data structure powers some of the most ubiquitous features in modern software, from search autocomplete to network routing. By following the step-by-step approach in this tutorial, you now have a complete, tested Trie implementation along with the knowledge of how to optimize it for different scenarios, extend it with deletion support, and avoid common mistakes. Whether you are preparing for a coding interview or building a real-world application that needs fast prefix lookups, the Trie is an indispensable tool in your algorithmic toolkit.

— Ad —

Google AdSense will appear here after approval

← Back to all articles