Implement Trie (Prefix Tree) Data Structure in Java

A Trie (Prefix Tree) is a tree-like data structure optimized for storing and retrieving strings with common prefixes. This implementation provides efficient insertion, search, and prefix matching operations with O(M) time complexity where M is the length of the word.

1. Problem Restatement

Clarifying questions to ask the interviewer:

1. What is the expected range of operations (insert, search, startsWith) that we'll be handling?
2. Are we allowed to modify the input trie structure during operations, or should it remain immutable?
3. Should we handle only lowercase letters (a-z) or also support other characters?
4. Do we need to consider case sensitivity in our implementation?
5. What's the target time complexity for each operation (O(1) or O(prefix length))?

Restatement of the problem: Implement a Trie (prefix tree) data structure that supports three core operations:

1. Insert a word into the trie
2. Search for an exact word in the trie
3. Check if any word in the trie starts with a given prefix

Input/Output specifications:

• Input: Words (strings) containing only lowercase letters a-z
• Operations: insert(String word), search(String word), startsWith(String prefix)
• Output: Boolean for search and startsWith operations

Constraints and assumptions:

• All words consist only of lowercase letters (a-z)
• Word lengths are reasonable (typically < 1000 characters)
• Operations will be called on a valid, initialized trie
• We can modify the trie structure during insert operations

Key insight: A Trie is a tree-like data structure where each node represents a character, and paths from root to leaves represent complete words. The structure naturally supports prefix-based operations.

Hidden structure: Each node maintains a mapping of child characters to child nodes, with an end-of-word marker to distinguish between prefixes and complete words. The tree structure enables efficient prefix matching by traversing character-by-character down the tree.

2. Pattern Recognition

Identifying the most likely pattern or patterns

The most likely pattern is Trie (Prefix Tree) construction and operations. This is a classic data structure problem that involves building a tree-like structure to efficiently store and retrieve strings based on their prefixes.

Wording that signals the pattern

• "Trie (Prefix Tree)" - direct signal of the data structure
• "prefix" - indicates we need to handle string prefix operations efficiently
• "insert", "search", "startsWith" - standard Trie operations that suggest we're building a tree where each node represents a character and paths represent strings

Why this problem fits that pattern

This problem fits the Trie pattern because:

1. We're dealing with string operations where prefix matching is essential
2. The problem asks for efficient insertion, search, and prefix checking operations
3. The structure naturally builds a tree where each path from root to leaf represents a complete word
4. The use case matches classic Trie applications like autocomplete, spell checkers, and IP routing

Core insight needed to solve it

The core insight is that a Trie node should store:

• A flag indicating if the node represents the end of a word
• An array/map of child nodes (one for each possible character)
• The ability to traverse the tree efficiently by character to build/lookup prefixes

The key optimization is using a fixed-size array for children (26 slots for lowercase letters) or a HashMap for more flexibility, allowing O(1) character-to-node mapping during traversal operations.

This pattern enables us to solve all three required operations in O(m) time where m is the length of the word/prefix being processed, making it much more efficient than naive string array approaches.

3. Approach Evolution

My first instinct is to use a nested HashMap or a simple array-based structure to represent the trie nodes, where each node stores references to its child nodes and a flag indicating if it's the end of a word. But that approach would be O(n²) in worst case for insertion and search operations because we'd need to traverse the entire prefix path for each operation, and in the worst case, we might end up checking every character of every word.

Let me think about how to optimize this. I notice that if I maintain a proper node structure with explicit children mapping and word-end markers, I can reduce it to O(L) where L is the length of the word being processed. The key insight is that we don't need to store entire words at each node, just the character mappings and a terminal flag.

The core optimization is to use a dedicated TrieNode class that stores:

1. A map of characters to child nodes (for flexibility with any alphabet)
2. A boolean flag indicating end-of-word
3. This allows us to build the trie incrementally and search efficiently by traversing character-by-character

I'm going to use a HashMap approach for the children because it's more scalable and handles any character set gracefully, rather than assuming ASCII or fixed-size arrays. The structure will allow me to build the trie in O(L) time per word where L is the word length, and search operations will also be O(L).

The invariant I must maintain is that each path from root to any node represents a valid prefix, and the terminal flag marks where words end. When inserting a new word, I ensure that all necessary intermediate nodes are created and the terminal flag is set correctly.

This approach gives us clean separation of concerns: insertion, search, and prefix checking all use the same traversal logic but with different terminal conditions.

The final approach will be:

1. Create a TrieNode class with HashMap of children and end-of-word flag
2. For insertion: traverse existing path, create new nodes as needed, mark end-of-word
3. For search: traverse path exactly matching the word, verify terminal flag
4. For prefix check: traverse path matching prefix, return true if we reach end of prefix

This is optimal for the given constraints and provides clean O(L) operations where L is the word/prefix length.

4. Data Structure Choice

For implementing a Trie (Prefix Tree), I'll use a combination of data structures:

Primary Data Structures:

1. Node-based Tree Structure - Each Trie node represents a character and contains:

   • An array of child pointers (size 26 for lowercase English letters)
   • A boolean flag to mark end of word
   • Optional: word storage for efficiency

2. Character Array/HashMap - For storing child nodes:

   • Using array of size 26 for fixed lowercase alphabet (more efficient than HashMap)
   • Each index corresponds to a letter (0='a', 1='b', ..., 25='z')

Why This Choice:

The array-based approach is optimal for this specific constraint (lowercase English letters only) because:

• O(1) access time for each child node
• No hash collision overhead
• Memory efficient with predictable size (26 * pointer size)
• Simpler to implement and debug than HashMap

Space Trade-off Consideration: While a HashMap would allow for dynamic character sets and handle Unicode, it introduces:

• Hash computation overhead
• Collision handling complexity
• Higher memory overhead due to hash table structure

For the standard problem constraints (lowercase letters only), the fixed-size array is superior in both time and space efficiency.

Helper Structure:

• Simple Node class with children[] array and isEnd boolean flag
• No additional helper structures needed for core operations

This choice enables O(1) character lookup and maintains clean, readable code while achieving optimal performance for the given constraints.

5. Algorithm Walkthrough

Let me trace through the algorithm using the example operations: insert("apple"), search("apple"), startsWith("app").

Step 1: Insert "apple"

• Start at root node
• For 'a': Create new node if doesn't exist, move to it
• For 'p': Create new node if doesn't exist, move to it
• For 'p': Create new node if doesn't exist, move to it
• For 'l': Create new node if doesn't exist, move to it
• For 'e': Create new node if doesn't exist, mark end of word

Step 2: Search "apple"

• Start at root node
• For 'a': Navigate to existing node
• For 'p': Navigate to existing node
• For 'p': Navigate to existing node
• For 'l': Navigate to existing node
• For 'e': Navigate to existing node
• Check if end marker is set - return true

Step 3: StartsWith "app"

• Start at root node
• For 'a': Navigate to existing node
• For 'p': Navigate to existing node
• For 'p': Navigate to existing node
• Check if we can traverse the path - return true

The algorithm maintains the invariant that each node in the trie represents a character, and the path from root to any node forms a valid prefix. When we insert a word, we ensure all intermediate nodes exist and mark the final node as end of word. When searching, we verify the complete path exists and the target is marked as a complete word. For prefix search, we only verify the path exists.

public class TrieNode {
    private static final int ALPHABET_SIZE = 26;
    private TrieNode[] children;
    private boolean isEndOfWord;

    public TrieNode() {
        children = new TrieNode[ALPHABET_SIZE];
        isEndOfWord = false;
    }

    public boolean containsKey(char ch) {
        return children[ch - 'a'] != null;
    }

    public TrieNode get(char ch) {
        return children[ch - 'a'];
    }

    public void put(char ch, TrieNode node) {
        children[ch - 'a'] = node;
    }

    public void setEnd() {
        isEndOfWord = true;
    }

    public boolean isEnd() {
        return isEndOfWord;
    }
}

public class Trie {
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            if (!node.containsKey(ch)) {
                node.put(ch, new TrieNode());
            }
            node = node.get(ch);
        }
        node.setEnd();
    }

    public boolean search(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            if (!node.containsKey(ch)) {
                return false;
            }
            node = node.get(ch);
        }
        return node.isEnd();
    }

    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            if (!node.containsKey(ch)) {
                return false;
            }
            node = node.get(ch);
        }
        return true;
    }
}

6. Java Solution

class TrieNode {
    private static final int ALPHABET_SIZE = 26;
    private TrieNode[] children;
    private boolean isEndOfWord;
    
    public TrieNode() {
        children = new TrieNode[ALPHABET_SIZE];
        isEndOfWord = false;
    }
    
    public boolean containsKey(char ch) {
        return children[ch - 'a'] != null;
    }
    
    public TrieNode get(char ch) {
        return children[ch - 'a'];
    }
    
    public void put(char ch, TrieNode node) {
        children[ch - 'a'] = node;
    }
    
    public void setEnd() {
        isEndOfWord = true;
    }
    
    public boolean isEnd() {
        return isEndOfWord;
    }
}

public class Trie {
    private TrieNode root;
    
    public Trie() {
        root = new TrieNode();
    }
    
    public void insert(String word) {
        TrieNode current = root;
        
        // Traverse the word character by character
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            
            // If current character is not present, create a new node
            if (!current.containsKey(ch)) {
                current.put(ch, new TrieNode());
            }
            
            // Move to the next node
            current = current.get(ch);
        }
        
        // Mark the end of the word
        current.setEnd();
    }
    
    public boolean search(String word) {
        TrieNode current = root;
        
        // Traverse the word character by character
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            
            // If character is not found, word doesn't exist
            if (!current.containsKey(ch)) {
                return false;
            }
            
            // Move to the next node
            current = current.get(ch);
        }
        
        // Return true only if we reached the end of a word
        return current.isEnd();
    }
    
    public boolean startsWith(String prefix) {
        TrieNode current = root;
        
        // Traverse the prefix character by character
        for (int i = 0; i < prefix.length(); i++) {
            char ch = prefix.charAt(i);
            
            // If prefix character is not found, prefix doesn't exist
            if (!current.containsKey(ch)) {
                return false;
            }
            
            // Move to the next node
            current = current.get(ch);
        }
        
        // If we can traverse the entire prefix, it exists
        return true;
    }
}

The solution implements a Trie (prefix tree) data structure with three core operations:

1. Insert - builds the trie by creating nodes for each character
2. Search - checks if a complete word exists in the trie
3. StartsWith - checks if any word in the trie starts with the given prefix

Each TrieNode maintains an array of 26 pointers (for lowercase English letters) and a boolean flag to mark end-of-word. The implementation handles all edge cases including empty strings and ensures proper traversal through the tree structure.

7. Complexity Analysis

Time Complexity:

• Insert operation: O(m) where m is the length of the word being inserted
• Search operation: O(m) where m is the length of the word being searched
• StartsWith operation: O(p) where p is the length of the prefix being searched

Each operation processes each character of the input string exactly once, making all operations linear in the length of the input string.

Space Complexity:

• O(ALPHABET_SIZE × N × M) where:
  • ALPHABET_SIZE = 26 (for lowercase English letters)
  • N = number of words stored in the trie
  • M = average length of words

In the worst case, every word has a unique path in the trie, requiring space for each character node. The space complexity is proportional to the total number of characters across all words stored.

Why these bounds are correct: The bounds are tight because:

1. Each character is processed exactly once per operation
2. In the worst case, each unique word requires a new node for each character in its path
3. The trie structure ensures no redundant storage - only necessary nodes are created

Optimality: This solution is optimal for the trie data structure. Any implementation that stores words with prefix relationships will require at least O(N×M) space to maintain the full trie structure, and operations cannot be better than O(M) time since we must check each character in the input string to verify membership or prefix existence.

The solution achieves the theoretical lower bounds for trie operations, making it near-optimal for this specific data structure and use case.

8. Edge Cases & Pitfalls

Empty and Minimal Inputs

• Empty trie operations: Inserting empty string should work, but queries for empty prefix should return valid results
• Single character operations: Inserting and searching single characters must handle correctly, especially when they're the only node
• Null inputs: The problem typically assumes valid strings, but defensive programming should handle null inputs gracefully

Duplicate Values and Repeated States

• Duplicate insertions: Inserting the same word multiple times should not create duplicate nodes or affect existing structure
• Overlapping prefixes: Inserting "app" and "application" should properly share the common prefix path
• Same prefix different lengths: Inserting "app" and "apple" should maintain proper tree structure with shared nodes

Boundary Conditions

• Empty string handling: Empty string should be treated as a valid word (if allowed) and handled in search operations
• Single node trees: When only one character exists, all operations must work correctly without null pointer exceptions
• Edge prefix matching: Searching for prefixes that don't exist should return false without throwing exceptions

Common Implementation Mistakes

• Missing wordEnd marker: Forgetting to mark the end of words can cause incorrect search results
• Incorrect node traversal: Not properly advancing through the trie structure when searching for prefixes or words
• Memory management issues: Not properly handling node cleanup in garbage-collected environments (though less relevant in Java)
• Null pointer exceptions: Failing to check for null child nodes before accessing them during traversal
• Incorrect prefix vs word search: Confusing whether to check for complete word or just prefix existence

Implementation-specific Pitfalls

• Node reuse: Reusing trie nodes incorrectly when clearing or resetting the structure
• Character range assumptions: Assuming only lowercase letters (a-z) when the problem allows any valid character set
• Memory leaks: In languages with manual memory management, but in Java the garbage collector handles this automatically
• Thread safety: If the trie is used in concurrent environments, proper synchronization must be considered (though not required for basic implementation)

9. Alternative Approaches

There are several alternative approaches to implementing a Trie, each with different trade-offs:

1. Array-based Trie Implementation Instead of using HashMap for children, we can use a fixed-size array of size 26 (for lowercase English letters). This approach is more memory-efficient for the specific use case of lowercase letters but less flexible for extended character sets.

class TrieNode {
    TrieNode[] children;
    boolean isEndOfWord;
    
    public TrieNode() {
        children = new TrieNode[26];
        isEndOfWord = false;
    }
}

2. Node-based Implementation with HashSet We could store children in a HashSet instead of HashMap, but this would require maintaining the order and would have O(n) lookup time where n is the number of children, making it less efficient than HashMap's O(1) average lookup.

3. Using Ternary Search Tree (TST) For very large datasets with many common prefixes, a TST can be more space-efficient than a standard Trie. However, it's more complex to implement and has higher constant factors for operations.

4. Compressed Trie (Patricia Trie) In a compressed trie, we merge nodes that have only one child. This reduces space usage but makes insertion and search operations more complex due to the need to split and merge nodes.

5. Using Java's Built-in Collections with Custom Logic Some developers might consider using a Map<Character, TrieNode> but implementing their own logic for handling the character-to-node mapping. This is essentially the same as our HashMap approach but with less clean code.

The main trade-offs between these alternatives:

• Space vs Time: Array-based approaches save space but lose flexibility; HashMap provides better performance for sparse character sets
• Complexity: TSTs and compressed tries are more complex to implement correctly
• Flexibility: HashMap approach is most adaptable to different character sets and use cases

Our chosen HashMap-based solution offers the best balance of simplicity, performance for typical use cases, and maintainability for production code. The array-based approach would be preferred only when we're certain about the character set and want to minimize memory overhead.

10. Test Cases

// Test case 1: Basic functionality with insert and search
Trie trie = new Trie();
trie.insert("apple");
assertTrue(trie.search("apple"));   // returns true
assertFalse(trie.search("app"));    // returns false
assertTrue(trie.startsWith("app")); // returns true

// Test case 2: Insert same word multiple times
trie.insert("app");
trie.insert("app");
assertTrue(trie.search("app"));     // returns true
assertTrue(trie.startsWith("app")); // returns true

// Test case 3: Empty string handling
trie.insert("");
assertTrue(trie.search(""));        // returns true
assertTrue(trie.startsWith(""));   // returns true

// Test case 4: Complex prefix matching
trie.insert("apples");
trie.insert("application");
assertTrue(trie.search("apples"));     // returns true
assertTrue(trie.search("application")); // returns true
assertFalse(trie.search("apply"));      // returns false
assertTrue(trie.startsWith("app"));    // returns true
assertTrue(trie.startsWith("appl"));   // returns true

// Test case 5: Single character words
trie.insert("a");
assertTrue(trie.search("a"));         // returns true
assertTrue(trie.startsWith("a"));     // returns true
assertFalse(trie.search("ab"));       // returns false
assertTrue(trie.startsWith("ab"));    // returns false

// Test case 6: Large prefix overlap
trie.insert("abc");
trie.insert("abcd");
trie.insert("abcde");
assertTrue(trie.search("abc"));       // returns true
assertTrue(trie.search("abcd"));      // returns true
assertTrue(trie.search("abcde"));     // returns true
assertTrue(trie.startsWith("abc"));   // returns true
assertTrue(trie.startsWith("abcd"));  // returns true

The test cases cover fundamental functionality including basic insert/search operations, handling of duplicate entries, empty strings, complex prefix scenarios with overlapping characters, and boundary conditions with single character words. These cases ensure the implementation correctly handles various edge cases that commonly trip up candidates during interviews.

11. Final Summary

The core idea is to implement a Trie data structure using a tree-like structure where each node represents a character and tracks word endings. This solution efficiently supports prefix-based operations by leveraging the tree structure to avoid redundant string comparisons, making it optimal for autocomplete and word search scenarios where prefix matching is frequent.

The solution works by maintaining a tree of character nodes with end-of-word markers, allowing O(m) time complexity for search operations where m is the length of the word/prefix. The efficient use of object references and the tree structure ensures that common prefixes are shared, minimizing memory overhead while maximizing lookup performance.

What makes this solution efficient is the careful balance between memory usage (each node only stores necessary information) and lookup speed (no redundant character comparisons), making it ideal for applications requiring fast prefix-based searching like autocomplete features or spell checkers.