Implement Trie ll | Full Trie with Insert, Count, and Erase

The Implement Trie II problem is an extension of the basic Implement Trie (Prefix Tree). In addition to inserting words and checking prefixes, this version requires advanced operations like:

1. Insert(word): Insert a word into the Trie.
2. countWordsEqualTo(word): Return the number of times a given word has been inserted.
3. countWordsStartingWith(prefix): Return the number of words that start with the given prefix.
4. erase(word): Remove one occurrence of the word from the Trie.

This extended Trie allows handling multiple occurrences of the same word and gives more control over word counts.

Here’s the [Problem Link] to begin with.

Example

Operations: 
insert("apple")
insert("apple")
countWordsEqualTo("apple")   → 2
countWordsStartingWith("app") → 2
erase("apple")
countWordsEqualTo("apple")   → 1

Explanation:

• "apple" was inserted twice.
• Both words share prefix "app".
• After erasing one "apple", only one remains.

Intuition and Approach

To implement Implement Trie II, we extend the traditional Trie by maintaining two counters at each node:

1. count_prefix:
   • Counts how many words pass through this node.
   • Useful for countWordsStartingWith.
2. end_count:
   • Counts how many words end at this node.
   • Useful for countWordsEqualTo.

Operations

1. Insert(word):
   • Traverse each character.
   • If the character path doesn’t exist, create a new node.
   • Increment count_prefix at every node visited.
   • At the final node, increment end_count.
2. countWordsEqualTo(word):
   • Traverse characters of the word.
   • If a character is missing, return 0.
   • At the end, return end_count of the final node.
3. countWordsStartingWith(prefix):
   • Traverse prefix characters.
   • If a character is missing, return 0.
   • Return count_prefix of the final prefix node.
4. erase(word):
   • Traverse the word’s path.
   • Decrease count_prefix at each node.
   • Decrease end_count at the final node.

Code Implementation

from os import *
from sys import *
from collections import *
from math import *

class TrieNode:
    def __init__(self):
        self.children = {}
        self.count_prefix = 0
        self.end_count = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
            node.count_prefix += 1
        node.end_count += 1

    def countWordsEqualTo(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                return 0
            node = node.children[ch]
        return node.end_count

    def countWordsStartingWith(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                return 0
            node = node.children[ch]
        return node.count_prefix

    def erase(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                return
            node = node.children[ch]
            node.count_prefix -= 1
        node.end_count -= 1

Code Explanation

• TrieNode:
  • children: stores next characters.
  • count_prefix: counts how many words include this node in their path.
  • end_count: counts how many words end here.
• insert: Updates prefix counts along the path and increments end count.
• countWordsEqualTo: Returns how many words exactly match.
• countWordsStartingWith: Returns how many words share a given prefix.
• erase: Decreases prefix counts and word end counts to remove one occurrence.

Dry Run

Operations:

insert("apple")
insert("apple")
countWordsEqualTo("apple")
countWordsStartingWith("app")
erase("apple")
countWordsEqualTo("apple")

Steps:

• Insert "apple" twice → end_count("e") = 2, prefix counts updated along path.
• countWordsEqualTo("apple") → 2.
• countWordsStartingWith("app") → 2.
• Erase "apple" → end_count("e") = 1, prefix counts decremented.
• countWordsEqualTo("apple") → 1.

Time and Space Complexity

• Insert: O(L), where L = length of word.
• countWordsEqualTo: O(L).
• countWordsStartingWith: O(P), where P = prefix length.
• erase: O(L).
• Space Complexity: O(N * L), where N = number of words, L = average length of words.

Conclusion

The Implement Trie II problem enhances the basic Trie by adding advanced operations like counting exact words, counting prefixes, and erasing occurrences.

• Supports efficient insert, search, count, and erase.
• Useful in text editors, autocomplete, spell checkers, and dictionaries.
• Runs in O(L) per operation with scalable memory usage.

This makes it one of the most powerful variations of Trie for interview preparation.

For any changes to the article, kindly email at code@codeanddebug.in or contact us at +91-9712928220.