Implement Trie II (Prefix Tree)

Problem Statement

You need to implement a trie (prefix tree), a specialized tree data structure used to efficiently store and query strings. This is useful in applications like autocomplete, spell checkers, and string dictionaries.

Your Trie class should support the following methods:

• Trie(): Initializes the trie object.
• void insert(String word): Inserts the string word into the trie. If the word already exists, its instance count increases.
• int countWordsEqualTo(String word): Returns how many times the exact word exists in the trie.
• int countWordsStartingWith(String prefix): Returns how many words in the trie start with the given prefix.
• void erase(String word): Removes one occurrence of word from the trie. If the word exists multiple times, only one instance is erased.

Examples

Example 1

Input:

["Trie", "insert", "insert", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsStartingWith"]
[[], ["apple"], ["apple"], ["apple"], ["app"], ["apple"], ["apple"], ["app"], ["apple"], ["app"]]

Output:

[null, null, null, 2, 2, null, 1, 1, null, 0]

Explanation:

Trie trie = new Trie();
trie.insert("apple");                      // Adds "apple"
trie.insert("apple");                      // Adds another "apple"
trie.countWordsEqualTo("apple");           // Returns 2
trie.countWordsStartingWith("app");        // Returns 2
trie.erase("apple");                       // Removes one "apple"
trie.countWordsEqualTo("apple");           // Returns 1
trie.countWordsStartingWith("app");        // Returns 1
trie.erase("apple");                       // Removes the last "apple"
trie.countWordsStartingWith("app");        // Returns 0

Constraints

• 1 <= word.length, prefix.length <= 2000
• word and prefix consist only of lowercase English letters.
• At most 3 * 10^4 calls will be made to all operations combined.
• It is guaranteed that erase(word) is only called when word exists in the trie.

Approach and Intuition

The trie data structure supports all required operations efficiently by leveraging character-wise traversal and storage:

1. Initialization:

   • The root node starts empty. Each node stores:

     • A map/array of child nodes (one for each character)
     • A count of words ending at that node
     • A count of words passing through that node (for prefix counting)

2. Insertion:

   • Traverse the trie character by character.

   • For each character:

     • Create a node if it doesn't exist.
     • Increment prefix count.

   • At the final node, increment the end count.

3. Counting Exact Matches:

   • Traverse through characters in word.
   • If traversal completes, return the node’s end count; otherwise, return 0.

4. Counting Prefix Matches:

   • Traverse through characters in prefix.
   • Return the prefix count of the last node if traversal succeeds; otherwise, return 0.

5. Erase:

   • Traverse down the trie as in insertion.
   • Decrement the prefix count along the path.
   • Decrement the end count at the final node.

This design ensures all operations run in O(L) time, where L is the length of the word or prefix, making it efficient even with tens of thousands of operations.

Solutions

class TrieNode {
public:
    TrieNode* children[26];
    int terminalCount = 0;
    int prefixCount = 0;
};
    
class Trie {
public:
    TrieNode* root;
    
    Trie() {
        root = new TrieNode();
    }
    
    void addWord(string word) {
        TrieNode* current = root;
        for (char& c : word) {
            int idx = c - 'a';
            if (!current->children[idx]) {
                current->children[idx] = new TrieNode();
            }
            current = current->children[idx];
            current->prefixCount++;
        }
        current->terminalCount++;
    }
    
    int countExactWords(string word) {
        TrieNode* current = root;
        for (char& c : word) {
            int idx = c - 'a';
            if (!current->children[idx]) {
                return 0;
            }
            current = current->children[idx];
        }
        return current->terminalCount;
    }
    
    int countPrefixWords(string prefix) {
        TrieNode* current = root;
        for (char& c : prefix) {
            int idx = c - 'a';
            if (!current->children[idx]) {
                return 0;
            }
            current = current->children[idx];
        }
        return current->prefixCount;
    }
    
    void removeWord(string word) {
        TrieNode* current = root;
        for (char& c : word) {
            int idx = c - 'a';
            current = current->children[idx];
            current->prefixCount--;
        }
        current->terminalCount--;
    }
};