Pular para o conteúdo principal

Implement Trie (Prefix Tree)

This question is asked by Microsoft. Implement a trie class that supports insertion and search functionalities.

Note: You may assume only lowercase alphabetical characters will be added to your trie.

Example(s)​

Example 1:

Trie trie = new Trie()
trie.insert("programming");
trie.search("computer") // returns false
trie.search("programming") // returns true

Example 2:

Trie trie = new Trie()
trie.insert("apple");
trie.search("app") // returns false (not a complete word)
trie.search("apple") // returns true
trie.insert("app");
trie.search("app") // returns true

Example 3:

Trie trie = new Trie()
trie.insert("cat");
trie.insert("car");
trie.insert("dog");
trie.search("cat") // returns true
trie.search("car") // returns true
trie.search("dog") // returns true
trie.search("ca") // returns false

Solution​

The solution uses a Trie node structure:

  1. TrieNode class: Each node has children (26 for lowercase letters) and an end-of-word flag
  2. Insert: Traverse/create path for each character, mark end of word
  3. Search: Traverse path, check if word exists and is marked as complete

JavaScript Solution - Trie Implementation

/**
* Trie Node class
*/
class TrieNode {
constructor() {
  this.children = new Map(); // Map<char, TrieNode>
  this.isEndOfWord = false; // Marks end of a word
}
}

/**
* Trie class
*/
class Trie {
constructor() {
  this.root = new TrieNode();
}

/**
 * Insert a word into the trie
 * @param {string} word - Word to insert
 * @return {void}
 */
insert(word) {
  let current = this.root;
  
  for (const char of word) {
    // If character doesn't exist, create new node
    if (!current.children.has(char)) {
      current.children.set(char, new TrieNode());
    }
    // Move to next node
    current = current.children.get(char);
  }
  
  // Mark end of word
  current.isEndOfWord = true;
}

/**
 * Search for a word in the trie
 * @param {string} word - Word to search
 * @return {boolean} - True if word exists
 */
search(word) {
  let current = this.root;
  
  for (const char of word) {
    // If character doesn't exist, word not found
    if (!current.children.has(char)) {
      return false;
    }
    // Move to next node
    current = current.children.get(char);
  }
  
  // Check if it's a complete word (not just a prefix)
  return current.isEndOfWord;
}

/**
 * Check if there's any word starting with the given prefix
 * @param {string} prefix - Prefix to check
 * @return {boolean} - True if prefix exists
 */
startsWith(prefix) {
  let current = this.root;
  
  for (const char of prefix) {
    if (!current.children.has(char)) {
      return false;
    }
    current = current.children.get(char);
  }
  
  return true; // Prefix exists (may or may not be complete word)
}
}

// Test cases
const trie = new Trie();
trie.insert("programming");
console.log('Search "computer":', trie.search("computer")); // false
console.log('Search "programming":', trie.search("programming")); // true
console.log('StartsWith "prog":', trie.startsWith("prog")); // true
Output:
Click "Run Code" to execute the code and see the results.

Alternative Solution (Array-based Children)​

Here's a version using arrays instead of maps:

/**
 * Trie Node with array-based children
 */
class TrieNodeArray {
  constructor() {
    this.children = new Array(26).fill(null); // 26 lowercase letters
    this.isEndOfWord = false;
  }
}

class TrieArray {
  constructor() {
    this.root = new TrieNodeArray();
  }
  
  insert(word) {
    let current = this.root;
    
    for (const char of word) {
      const index = char.charCodeAt(0) - 'a'.charCodeAt(0);
      
      if (!current.children[index]) {
        current.children[index] = new TrieNodeArray();
      }
      current = current.children[index];
    }
    
    current.isEndOfWord = true;
  }
  
  search(word) {
    let current = this.root;
    
    for (const char of word) {
      const index = char.charCodeAt(0) - 'a'.charCodeAt(0);
      
      if (!current.children[index]) {
        return false;
      }
      current = current.children[index];
    }
    
    return current.isEndOfWord;
  }
}

Complexity​

  • Time Complexity:
    • Insert: O(m) - Where m is the length of the word
    • Search: O(m) - Where m is the length of the word
  • Space Complexity: O(ALPHABET_SIZE × N × M) - Where N is the number of words and M is the average length. In practice, it's O(total characters in all words).

Approach​

The solution uses a Trie (Prefix Tree) data structure:

  1. TrieNode structure:

    • Each node has children (map or array)
    • Each node has an isEndOfWord flag

  2. Insert operation:

    • Traverse/create path for each character
    • Create new nodes if needed
    • Mark end of word at the last node

  3. Search operation:

    • Traverse path character by character
    • Return false if path doesn't exist
    • Return true only if path exists AND is marked as end of word

Key Insights​

  • Trie structure: Tree where each path represents a word
  • Shared prefixes: Words with common prefixes share nodes
  • End-of-word flag: Distinguishes complete words from prefixes
  • O(m) operations: Both insert and search are O(m) where m is word length
  • Space efficient: Shares common prefixes among words

Step-by-Step Example​

Let's trace through inserting "programming":

Initial: root = TrieNode()

Insert "programming":
  current = root
  
  'p': children['p'] doesn't exist → create node
    current = children['p']
  
  'r': children['r'] doesn't exist → create node
    current = children['r']
  
  'o': children['o'] doesn't exist → create node
    current = children['o']
  
  ... (continue for all characters)
  
  'g': children['g'] doesn't exist → create node
    current = children['g']
    current.isEndOfWord = true

Trie structure:
  root
   |
   p
   |
   r
   |
   o
   |
   g
   |
   r
   |
   a
   |
   m
   |
   m
   |
   i
   |
   n
   |
   g (isEndOfWord = true)

Visual Representation​

Example: Insert "programming", "pro", "code"

Trie structure:
        root
       /    \
      p      c
      |      |
      r      o
      |      |
      o      d
     / \     |
    g*  m    e*
    |   |
    ... i
        |
        n
        |
        g*

* = isEndOfWord = true

Search "programming":
  Follow path: p → r → o → g → r → a → m → m → i → n → g
  Reached node with isEndOfWord = true → return true

Search "pro":
  Follow path: p → r → o
  Node exists but isEndOfWord = false → return false

Edge Cases​

  • Empty word: Can insert empty string (mark root as end of word)
  • Single character: Works correctly
  • Duplicate words: Inserting same word twice is fine
  • Prefix of existing word: Works correctly (e.g., "app" and "apple")
  • Word longer than existing: Works correctly

Important Notes​

  • End-of-word flag: Critical to distinguish complete words from prefixes
  • Shared prefixes: Words share common prefix nodes
  • Lowercase only: Problem states only lowercase letters
  • O(m) operations: Both insert and search are linear in word length
  • Space efficient: More efficient than storing all words separately

Why Trie is Efficient​

Key advantages:

  • Prefix sharing: Words with common prefixes share nodes
  • Fast lookup: O(m) time for search/insert
  • Space efficient: Only stores unique characters
  • Prefix queries: Easy to check if prefix exists

Example:

Words: "app", "apple", "application"
Shared prefix "app" is stored once, not three times
  • Add and Search Words: Trie with wildcard support
  • Word Search II: Use trie for word search
  • Longest Word in Dictionary: Different trie problem
  • Replace Words: Different trie application

Takeaways​

  • Trie structure efficiently stores and searches words
  • End-of-word flag distinguishes complete words from prefixes
  • O(m) operations for both insert and search
  • Shared prefixes make it space efficient
  • Tree structure with character edges