Problem Description

This problem asks you to implement a Trie (also known as a prefix tree) data structure with enhanced functionality. A Trie is a tree-based data structure that efficiently stores and retrieves strings, commonly used in applications like autocomplete and spell checkers.

You need to implement a [Trie](/problems/trie_intro) class with the following methods:

1. [Trie](/problems/trie_intro)(): Constructor that initializes an empty trie object.

2. insert(word): Adds a string word to the trie. If the same word is inserted multiple times, the trie should keep track of how many times it has been inserted.

3. countWordsEqualTo(word): Returns the count of how many times the exact string word exists in the trie. If the word has been inserted 3 times, this should return 3.

4. countWordsStartingWith(prefix): Returns the total count of all strings in the trie that begin with the given prefix. For example, if the trie contains "apple" (inserted twice) and "application" (inserted once), then countWordsStartingWith("app") should return 3.

5. erase(word): Removes one instance of the string word from the trie. This assumes the word exists in the trie before calling erase.

The key difference from a standard Trie implementation is that this version needs to:

• Track multiple occurrences of the same word (a word can be inserted multiple times)
• Count words with specific prefixes considering their multiplicity
• Support deletion of words while maintaining accurate counts

Each node in the solution maintains two counters:

• v: The number of times a word ends at this node
• pv: The total number of words that pass through this node (prefix count)

Intuition

The core challenge here is that we need to track not just whether words exist, but how many times each word appears and how many words share common prefixes. This means a standard Trie where nodes simply mark word endings isn't sufficient.

Think about what happens when we insert "apple" three times and "application" twice. When we ask for words starting with "app", we need to return 5 (not 2). This tells us we need to maintain running counts as we traverse the Trie.

The key insight is to maintain two separate counters at each node:

1. Word count (v): How many words end exactly at this node. This directly answers countWordsEqualTo().

2. Prefix count (pv): How many words pass through this node on their way to completion. This is the sum of all words that have this node's path as their prefix, which directly answers countWordsStartingWith().

Why do we need both? Consider inserting "app" and "apple":

• At node 'p' (third character), pv = 2 (both words pass through) but v = 1 (only "app" ends here)
• At node 'l', pv = 1 and v = 0 (only "apple" passes through, doesn't end)
• At node 'e', pv = 1 and v = 1 (only "apple" passes through and ends)

When we insert a word, we increment pv for every node along the path (since this word uses all these nodes as prefixes), but only increment v at the final node where the word ends.

For deletion, we simply reverse this process: decrement pv along the entire path and decrement v at the word's ending node. We don't need to physically remove nodes because the counters tell us everything - a node with pv = 0 is effectively "not used" even if it still exists in memory.

This design elegantly handles multiple insertions of the same word - inserting "apple" three times means the final 'e' node will have v = 3, and all nodes along the path will have their pv increased by 3 total.

Learn more about Trie patterns.

Solution Approach

The implementation uses an array-based Trie where each node contains:

• children: An array of 26 pointers (for lowercase letters a-z)
• v: Counter for words ending at this node
• pv: Counter for words passing through this node (prefix count)

1. Insert String

Starting from the root, we traverse the Trie character by character:

def insert(self, word: str) -> None:
    node = self
    for c in word:
        idx = ord(c) - ord('a')  # Convert char to index (0-25)
        if node.children[idx] is None:
            node.children[idx] = [Trie](/problems/trie_intro)()  # Create new node if needed
        node = node.children[idx]
        node.pv += 1  # Increment prefix count for this node
    node.v += 1  # Increment word count at final node

For each character, we:

1. Calculate its index using ord(c) - ord('a') to map 'a'→0, 'b'→1, ..., 'z'→25
2. Create a child node if it doesn't exist
3. Move to the child node and increment its pv (this word uses this node as prefix)
4. After processing all characters, increment v at the final node

Time complexity: O(n) where n is the word length.

2. Search Helper Function

The search method is a utility function that navigates to a specific node:

def search(self, word):
    node = self
    for c in word:
        idx = ord(c) - ord('a')
        if node.children[idx] is None:
            return None  # Path doesn't exist
        node = node.children[idx]
    return node  # Return the node where word/prefix ends

This returns the node where the given string ends, or None if the path doesn't exist.

Time complexity: O(n) where n is the word length.

3. Count Words Equal To

def countWordsEqualTo(self, word: str) -> int:
    node = self.search(word)
    return 0 if node is None else node.v

Uses the search helper to find the node where word ends, then returns its v value (number of times this exact word was inserted).

Time complexity: O(n) where n is the word length.

4. Count Words Starting With

def countWordsStartingWith(self, prefix: str) -> int:
    node = self.search(prefix)
    return 0 if node is None else node.pv

Uses the search helper to find the node where prefix ends, then returns its pv value (total count of all words passing through this node).

Time complexity: O(n) where n is the prefix length.

5. Erase String

def erase(self, word: str) -> None:
    node = self
    for c in word:
        idx = ord(c) - ord('a')
        node = node.children[idx]
        node.pv -= 1  # Decrement prefix count
    node.v -= 1  # Decrement word count at final node

Assumes the word exists (as per problem constraints). We traverse the path and:

1. Decrement pv for each node along the path
2. Decrement v at the final node

Note that we don't physically delete nodes even if their counters reach 0. This simplifies the implementation and doesn't affect correctness since we check counters, not node existence.

Time complexity: O(n) where n is the word length.

Example Walkthrough

If we insert "app" twice and "apple" once:

• After inserting "app" twice:
  • Nodes a, p, p have pv = 2
  • Last 'p' node has v = 2
• After inserting "apple":
  • Nodes a, p, p have pv = 3 (incremented by 1)
  • Nodes l, e have pv = 1
  • Node 'e' has v = 1
• countWordsStartingWith("app") returns 3 (from node 'p' with pv = 3)
• countWordsEqualTo("app") returns 2 (from node 'p' with v = 2)

Example Walkthrough

Let's walk through a concrete example to understand how the enhanced Trie works with its dual-counter system.

Scenario: We'll insert "cat" twice, "cats" once, and "dog" once, then perform various operations.

Step 1: Insert "cat" (first time)

Starting from root, traverse c → a → t
- At 'c': pv = 1 (one word passes through)
- At 'a': pv = 1 
- At 't': pv = 1, v = 1 (word ends here)

Step 2: Insert "cat" (second time)

Traverse the same path c → a → t
- At 'c': pv = 2 (two words pass through now)
- At 'a': pv = 2
- At 't': pv = 2, v = 2 (two instances of "cat" end here)

Step 3: Insert "cats"

Traverse c → a → t → s
- At 'c': pv = 3 (incremented, three words total pass through)
- At 'a': pv = 3
- At 't': pv = 3 (incremented, "cats" passes through here)
- At 's': pv = 1, v = 1 (only "cats" reaches and ends here)

Step 4: Insert "dog"

Create new path d → o → g
- At 'd': pv = 1
- At 'o': pv = 1
- At 'g': pv = 1, v = 1

Current Trie State:

root
├── c (pv=3)
│   └── a (pv=3)
│       └── t (pv=3, v=2)
│           └── s (pv=1, v=1)
└── d (pv=1)
    └── o (pv=1)
        └── g (pv=1, v=1)

Query Operations:

• countWordsEqualTo("cat") → Returns 2 (node 't' has v=2)
• countWordsStartingWith("ca") → Returns 3 (node 'a' has pv=3, counting both "cat" instances and "cats")
• countWordsStartingWith("cat") → Returns 3 (node 't' has pv=3)
• countWordsEqualTo("cats") → Returns 1 (node 's' has v=1)

Step 5: Erase "cat" (removes one instance)

Traverse c → a → t and decrement counters
- At 'c': pv = 2 (decremented from 3)
- At 'a': pv = 2
- At 't': pv = 2, v = 1 (decremented from v=2)
- Node 's' remains unchanged (pv=1, v=1)

After erasure:

• countWordsEqualTo("cat") → Returns 1 (one instance remains)
• countWordsStartingWith("ca") → Returns 2 (one "cat" + one "cats")

This example demonstrates how:

1. The pv counter tracks all words flowing through each node
2. The v counter tracks only words ending at that specific node
3. Multiple insertions of the same word increase both counters appropriately
4. Erasure decrements counters without removing nodes, maintaining the Trie structure

Solution Implementation

Time and Space Complexity

Time Complexity:

• insert(word): O(n) where n is the length of the word. The method traverses through each character of the word once, performing constant time operations (array access, arithmetic) at each step.

• countWordsEqualTo(word): O(n) where n is the length of the word. It calls the search method which traverses through the word once, then returns the count in constant time.

• countWordsStartingWith(prefix): O(n) where n is the length of the prefix. Similar to countWordsEqualTo, it uses the search method to traverse the prefix.

• erase(word): O(n) where n is the length of the word. The method traverses through each character of the word once, decrementing counters at each node.

• search(word): O(n) where n is the length of the word. It traverses through each character exactly once.

Space Complexity:

• For the Trie structure itself: O(ALPHABET_SIZE × N × L) in the worst case, where ALPHABET_SIZE = 26, N is the total number of unique words inserted, and L is the average length of words. However, in practice, due to prefix sharing, the space usage is often much less.

• For individual operations:

  • insert: O(n) in the worst case when creating new nodes for a word of length n
  • countWordsEqualTo, countWordsStartingWith, search: O(1) additional space (only using variables for traversal)
  • erase: O(1) additional space (only traversing existing nodes)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Not Handling Multiple Insertions of the Same Word

Pitfall: A common mistake is treating the Trie like a set where each word can only exist once. Developers might use a boolean flag (isEndOfWord) instead of a counter, which fails when the same word needs to be tracked multiple times.

Incorrect Implementation:

class Trie:
    def __init__(self):
        self.children = [None] * 26
        self.isEndOfWord = False  # Wrong! Can't track multiple insertions
        self.prefix_count = 0
  
    def insert(self, word):
        # ...
        node.isEndOfWord = True  # This doesn't count duplicates

Solution: Use integer counters (word_count and prefix_count) instead of boolean flags to track the frequency of words and prefixes.

2. Forgetting to Update Prefix Counts Along the Path

Pitfall: Only updating the counter at the final node without incrementing prefix_count for all nodes along the path. This breaks the countWordsStartingWith functionality.

Incorrect Implementation:

def insert(self, word):
    node = self
    for char in word:
        index = ord(char) - ord('a')
        if node.children[index] is None:
            node.children[index] = Trie()
        node = node.children[index]
        # Forgot to update node.prefix_count here!
    node.word_count += 1
    node.prefix_count += 1  # Wrong! Only updating at the end

Solution: Increment prefix_count for every node visited during insertion, not just the final node.

3. Physical Node Deletion in Erase Method

Pitfall: Attempting to physically remove nodes when erasing a word, which complicates the implementation and can break the tree structure if other words share the same prefix.

Incorrect Implementation:

def erase(self, word):
    def _erase_helper(node, word, index):
        if index == len(word):
            node.word_count -= 1
            # Trying to delete the node if empty
            return node.word_count == 0 and not any(node.children)
      
        char_index = ord(word[index]) - ord('a')
        should_delete = _erase_helper(node.children[char_index], word, index + 1)
      
        if should_delete:
            node.children[char_index] = None  # Dangerous! Other words might use this path
      
        node.prefix_count -= 1
        return node.prefix_count == 0 and not any(node.children)

Solution: Simply decrement the counters without physically removing nodes. Nodes with zero counts are effectively "empty" and don't affect the correctness of operations.

4. Not Checking for None in Search Before Accessing Properties

Pitfall: Directly accessing properties of the search result without checking if it's None first, leading to AttributeError.

Incorrect Implementation:

def countWordsEqualTo(self, word):
    node = self.search(word)
    return node.word_count  # AttributeError if node is None!

Solution: Always check if the search result is None before accessing its properties:

def countWordsEqualTo(self, word):
    node = self.search(word)
    return 0 if node is None else node.word_count

5. Assuming Word Exists in Erase Without Validation

Pitfall: While the problem states that erase assumes the word exists, in production code, not validating this can lead to negative counters or incorrect state.

Potential Issue:

# If someone calls erase("word") when "word" doesn't exist:
# - prefix_count values become negative
# - word_count becomes negative
# - Future operations give incorrect results

Solution for Production Code: Add validation to ensure the word exists before erasing:

def erase(self, word):
    # First check if word exists
    if self.countWordsEqualTo(word) == 0:
        return  # Or raise an exception
  
    # Then proceed with normal erase logic
    node = self
    for char in word:
        index = ord(char) - ord('a')
        node = node.children[index]
        node.prefix_count -= 1
    node.word_count -= 1