Trie-based Search and Autocomplete

1. Building an Autocomplete Feature

An autocomplete system suggests a list of complete words based on a prefix a user has typed. Our startsWith(prefix) method from the last lesson can tell us if any words have a given prefix, but it can't tell us what those words are. Let's fix that.

The core idea is simple:

1. Traverse the Trie to the node that represents the end of the input prefix.
2. From that node, perform a search (specifically, a Depth-First Search or DFS) to explore all possible downward paths.
3. Any path that ends at a node marked with isEndOfWord = true represents a valid suggestion.

Let's watch a short video that implements exactly this logic.

Implementation in Java

As you saw in the video, the implementation revolves around a recursive helper function. Let's formalize this. We'll add a method suggest(String prefix) to our Trie class.

First, the main suggest method finds the node corresponding to the end of the prefix.

// Inside your Trie class
public List<String> suggest(String prefix) {
    List<String> suggestions = new ArrayList<>();
    TrieNode currentNode = root; // Assuming 'root' is the start of your Trie

    // 1. Traverse to the end of the prefix
    for (char c : prefix.toCharArray()) {
        int index = c - 'a';
        if (currentNode.children[index] == null) {
            return suggestions; // Prefix not found, return empty list
        }
        currentNode = currentNode.children[index];
    }

    // 2. Call the recursive helper from the prefix's end node
    collectWords(currentNode, prefix, suggestions);
    
    return suggestions;
}

Next, we need the recursive helper method, collectWords. This method explores all paths from a given node.

// A private helper method inside your Trie class
private void collectWords(TrieNode node, String currentPrefix, List<String> suggestions) {
    // If the current node represents the end of a word, add it to suggestions
    if (node.isEndOfWord) {
        suggestions.add(currentPrefix);
    }

    // Recursively explore all children
    for (char c = 'a'; c <= 'z'; c++) {
        int index = c - 'a';
        if (node.children[index] != null) {
            // Append the character and recurse
            collectWords(node.children[index], currentPrefix + c, suggestions);
        }
    }
}

This pattern—traversing to a point and then launching a DFS—is a very common and powerful technique.

2. Word Search on a 2D Grid (The "Word Search II" Problem)

Now for a more complex and very popular interview problem. Imagine you are given a 2D grid of characters (like a Boggle board) and a dictionary of words. Your task is to find all the words from the dictionary that can be formed by connecting adjacent letters on the grid (horizontally or vertically).

Example:
Grid:
[['o', 'a', 'a', 'n'],
['e', 't', 'a', 'e'],
['i', 'h', 'k', 'r'],
['i', 'f', 'l', 'v']]

Dictionary: ["oath", "pea", "eat", "rain"]

Result: ["oath", "eat"]

The Naive Approach vs. The Trie Approach

A naive approach would be to iterate through each word in the dictionary and then, for each word, search the entire grid to see if it exists. This is incredibly inefficient. If your dictionary contains "apple", "apply", and "application", you would repeatedly search the grid for the common prefix "app".

This is where the Trie shines. We can flip the problem on its head: instead of checking each word against the grid, we can explore the grid and see what words we can form, using the Trie to guide our search.

The Trie-Optimized Algorithm:

1. Build Trie: Insert all words from the dictionary into a Trie.
2. DFS on Grid: Start a Depth-First Search (DFS) from every cell in the grid.
3. Prune with Trie: As you traverse the grid during the DFS, you build up a string of characters. At each step, check if this string exists as a prefix in the Trie.
   • If it is a valid prefix, continue the search to adjacent cells.
   • If it is not a valid prefix, there's no point in continuing down this path. Prune this branch of the search and backtrack immediately.
4. Collect Words: If your current path forms a complete word in the Trie (check the isEndOfWord flag), add it to your results list.

This video provides an excellent walkthrough of this exact problem and solution.

Java Implementation Insights

Let's break down the dfs backtracking function, as it's the heart of the solution. It combines grid traversal logic with Trie lookups.

// Main method outline
public List<String> findWords(char[][] board, String[] words) {
    Trie trie = new Trie();
    for (String word : words) {
        trie.insert(word);
    }

    Set<String> result = new HashSet<>(); // Use a Set to handle duplicates

    // Start DFS from every cell
    for (int i = 0; i < board.length; i++) {
        for (int j = 0; j < board[0].length; j++) {
            dfs(board, i, j, trie.getRoot(), result);
        }
    }
    return new ArrayList<>(result);
}

// DFS helper function
private void dfs(char[][] board, int i, int j, TrieNode node, Set<String> result) {
    // 1. Base cases for stopping the recursion
    if (i < 0 || i >= board.length || j < 0 || j >= board[0].length) {
        return; // Out of bounds
    }
    char c = board[i][j];
    if (c == '#' || node.children[c - 'a'] == null) {
        return; // Cell already visited or path is not a valid prefix in the Trie
    }

    // 2. Move to the next node in the Trie
    node = node.children[c - 'a'];

    // 3. Check if we found a word
    if (node.word != null) { // Using the video's optimization of storing the word in the node
        result.add(node.word);
        node.word = null; // De-duplicate: Avoid finding the same word again from the same path
    }

    // 4. Mark cell as visited and explore neighbors
    board[i][j] = '#'; // Make change
    
    dfs(board, i + 1, j, node, result); // Down
    dfs(board, i - 1, j, node, result); // Up
    dfs(board, i, j + 1, node, result); // Right
    dfs(board, i, j - 1, node, result); // Left

    // 5. Backtrack: un-mark the cell
    board[i][j] = c; // Undo change
}

Note: This code assumes a TrieNode structure that includes TrieNode[] children and String word (which is null if it's not an end-of-word node).

Conclusion

Today, you've seen how to elevate a basic Trie into a sophisticated tool for solving complex problems. We moved beyond simple insertion and search to implement full-fledged features that are staples of software engineering and interview questions.

Key Takeaways:

• Autocomplete is implemented by performing a DFS from the Trie node corresponding to the end of a given prefix.
• The Word Search II problem is a classic example of combining a data structure (Trie) with an algorithm (backtracking/DFS).
• The Trie acts as a "guide" for the DFS, allowing for aggressive pruning of the search space, which turns an otherwise intractable problem into an efficient one.
• Backtracking is a powerful pattern: make a change, recursively explore, and then undo the change.

Preview of the Next Lesson:
The Word Search II problem is a significant one, and mastering it is key. In the next lesson, we will solidify your understanding by focusing entirely on this pattern. We will solve problems involving finding all words from a dictionary on a 2D grid of characters using a trie and backtracking, looking at variations and further optimizations to ensure you're fully prepared to tackle it in any context.