Trie for Autocomplete and Prefix Matching

In our previous lesson, you built a fully functional Trie from the ground up, implementing the insert, search, and startsWith methods. We established that the startsWith operation is what makes a Trie uniquely powerful compared to a hash map. Today, we will harness that power for a practical and common application.

This lesson focuses on a direct application of the Trie's prefix-based nature: using a Trie to efficiently implement an autocomplete or prefix-matching feature. We will move beyond simply verifying if a prefix exists to retrieving all the words that begin with that prefix. This is the fundamental logic behind the suggestion features you see in search engines, code editors, and financial data terminals every day.

From Prefix Checking to Prefix Matching

In the last lesson, our startsWith method could answer a yes/no question: "Does any word in our dictionary begin with 'hel'?" Now, we want to answer a more useful question: "What are all the words in our dictionary that begin with 'hel'?"

The process can be broken down into two main steps:

1. Navigate to the Prefix Node: First, we traverse the Trie to find the node that represents the end of the given prefix (e.g., the node for 'l' in "hel").
2. Collect All Words Below: From that prefix node, we must explore its entire subtree to find all paths that lead to a complete word.

This diagram perfectly illustrates the concept. To find suggestions for the prefix "th", we first navigate the path t -> h. From the 'h' node, we then explore all branches below it, which lead us to the complete words "this" and "there".

The Algorithm: Navigate and Explore

Let's formalize the two steps.

1. Navigate to the Prefix Node

This step is identical to the _search_prefix helper method we developed in the previous lesson. We start at the root and iterate through the characters of the prefix. If at any point the path is broken (a child node is None), it means the prefix doesn't exist in our Trie, and we can immediately return an empty list of suggestions. If the traversal completes, we have our starting node for the next step.

2. Explore the Subtree with Depth-First Search (DFS)

Once we have the node corresponding to the end of the prefix, we need to find every valid word in the subtree rooted at that node. A recursive Depth-First Search is a natural and elegant way to accomplish this.

We can define a helper function, let's call it _collect_all_words, which takes a starting node and the prefix string that led to it. The logic would be:

• Check the current node: If the isEndOfWord flag of the current node is True, it means the path we've taken forms a complete word. We add this word to our list of results.
• Recurse on children: For each child of the current node, we make a recursive call. We pass the child node and the updated path (the current path plus the character corresponding to that child).

This process systematically explores every possible path downwards from the prefix node, collecting all valid words along the way.

Implementation in Practice

The following resources demonstrate how to implement this logic. We'll start with a clear article from GeeksforGeeks that provides code in several languages, and then supplement it with a video explanation.

The core of this algorithm is the recursive traversal. To see it in action, the following video provides an excellent animated walkthrough.

A Production-Oriented Perspective

Your experience as an intern has likely shown you that code is often viewed in the context of a larger system. The following resource frames our autocomplete problem within a narrative of building a production search system, which you might find insightful.

Complete Python Implementation

Let's integrate this new functionality into the Trie class from our previous lesson. Here is how you could add an autocomplete method.

# Assuming the TrieNode and Trie class from the previous lesson
class Trie:
    def __init__(self):
        self.root = TrieNode()




    # ... insert, _search_prefix, search, startsWith methods from last lesson ...
    
    def autocomplete(self, prefix: str) -> list[str]:
        """
        Returns all words in the trie that start with the given prefix.
        """



        # Step 1: Navigate to the node representing the end of the prefix.
        prefix_node = self._search_prefix(prefix)
        
        if not prefix_node:
            return [] # Prefix not found in the trie.




        # Step 2: From the prefix node, collect all words in the subtree.
        suggestions = []
        self._collect_all_words(prefix_node, prefix, suggestions)
        
        return suggestions

    def _collect_all_words(self, node: TrieNode, current_path: str, suggestions: list[str]):
        """
        A recursive helper to perform DFS from a given node.
        """



        # If the current node marks the end of a word, add it to our list.
        if node.isEndOfWord:
            suggestions.append(current_path)




        # Recursively call on all children, appending the character to the path.
        for i, child_node in enumerate(node.children):
            if child_node:
                child_char = chr(ord('a') + i)
                self._collect_all_words(child_node, current_path + child_char, suggestions)

In this implementation, autocomplete orchestrates the two-step process, while the recursive _collect_all_words helper does the heavy lifting of the DFS traversal. This separation of concerns makes the code clean and easy to follow.

Conclusion

In this lesson, you've leveled up your Trie skills from simple lookups to powerful prefix-matching. This is a classic interview pattern and a cornerstone of many real-world applications.

Key Takeaways:

• Implementing an autocomplete feature is a two-step process: navigate to the prefix node, then explore its subtree.
• The navigation step reuses the same logic as the startsWith operation.
• The exploration step is naturally implemented using a recursive Depth-First Search (DFS) that collects all words marked with isEndOfWord.
• The complete implementation combines a traversal method to find the prefix node with a recursive helper to collect all words from that point downward.

You now have a solid understanding of how a Trie's structure enables efficient prefix-based queries. This capability is not just for autocomplete; it's a powerful tool for optimizing algorithms that deal with sets of strings.

In our next lesson, we will explore one such optimization. You will learn how a Trie can dramatically improve the performance of a backtracking search for words on a 2D grid, a problem commonly known as Word Search II. The startsWith logic you've mastered will be the key to pruning the search space and avoiding dead-end paths.