Implement Trie II: Count Words, Prefixes and Erase

117.2k
0

Implement an advanced Trie that supports these operations:

  • insert(word) to add a word

  • countWordsEqualTo(word) to count how many times that exact word was inserted

  • countWordsStartingWith(prefix) to count how many stored words begin with that prefix

  • erase(word) to remove one occurrence of the word if it exists

Only lowercase English letters are considered in this version.

Example 1

Input: insert("apple"), insert("apple"), insert("app"), countWordsEqualTo("apple"), countWordsStartingWith("app"), erase("apple"), countWordsEqualTo("apple")

Output: 2, 3, 1

Explanation: "apple" is inserted twice and "app" once. So the exact count of "apple" becomes 2, and three stored words begin with "app". After erasing one "apple", its exact count becomes 1.

Example 2

Input: insert("bat"), insert("batch"), insert("bat"), countWordsEqualTo("bat"), countWordsStartingWith("bat"), erase("bat"), countWordsStartingWith("bat")

Output: 2, 3, 2

Explanation: The word "bat" appears twice, and both "bat" and "batch" contribute to the prefix "bat". After removing one occurrence of "bat", two stored words still begin with that prefix.

Approach

The turning point comes from asking what the basic Trie is missing. A normal Trie can tell whether a path exists, but it cannot answer questions like "How many times was apple inserted?" or "How many words currently start with app?" That means simple path traversal is no longer enough. Some extra information must be remembered while words are being inserted.

The useful condition is this: every time a word is inserted, it affects two kinds of places. First, it passes through several nodes while building its prefix path. Second, it ends at exactly one final node. Those are two different events, so they should be counted separately. Once that idea clicks, the logic becomes natural: store one counter for how many words pass through a node and another counter for how many words end there. Then exact-word count comes from the ending counter, prefix count comes from the pass-through counter, and erase simply reverses one earlier insertion.

Key Points

  • wordCount stores how many times the exact word ends at a node.

  • prefixCount stores how many inserted words pass through that node.

  • Duplicate insertions are allowed, so counters may become greater than 1.

  • erase(word) should do nothing if the word is not present.

  • This version performs logical deletion by decreasing counters. It does not physically free unused nodes.

Algorithm

  • Create a Trie node with 26 child references, one exact-word counter, and one prefix counter. This is needed because advanced queries ask two different things: how many words end here and how many words pass through here.

  • insert(word):

    • For insert(word), move character by character from the root. A missing node is created because the Trie must build the path before it can remember anything about that part of the word.

    • After moving to each child during insertion, increase its prefix counter. This is done because every visited node lies on the prefix path of the inserted word.

    • After the last character is processed, increase the exact-word counter of the final node. This step is necessary because only the last node should remember that a complete word ends there.

  • countWordsEqualTo(word):

    • For countWordsEqualTo(word), follow the exact path of the word. If the path breaks at any character, return 0 immediately because a missing path means that exact word was never stored fully.

    • If the full word path exists, return the exact-word counter of the last node. That counter is the direct answer because it records how many complete copies of that word were inserted.

  • countWordsStartingWith(prefix):

    • For countWordsStartingWith(prefix), again follow the path character by character. If the path breaks, return 0 because no stored word can continue from a prefix that does not exist.

    • If the full prefix path exists, return the prefix counter of the last prefix node. This works because that node has already counted how many words passed through it during insertion.

  • erase(word):

    • For erase(word), first check whether the word exists at least once. This safeguard is needed so the counters are not reduced for a word that was never present.

    • If the word exists, walk through the same path again and decrease the prefix counter at each visited node. This is done because one stored word is no longer using that prefix path.

    • After reaching the final node, decrease its exact-word counter. This step removes one full occurrence of the word while keeping any remaining duplicates intact.

Dry Run

Trie Implementation and Advanced Operations Dry Run

Trie Implementation and Advanced Operations Dry Run

Solution

#include <bits/stdc++.h>
using namespace std;
class TrieNode {
public:
vector<TrieNode*> children;
int wordCount;
int prefixCount;
// Stores child links and counters for one Trie node.
TrieNode() {
children.assign(26, nullptr);
wordCount = 0;
prefixCount = 0;
}
};
class Solution {
private:
TrieNode* root;
// Converts a lowercase character into its array position.
int getIndex(char ch) {
return ch - 'a';
}
public:
// Creates the Trie with an empty root node.
Solution() {
root = new TrieNode();
}
/*
Inserts a word and updates exact-word and prefix counters.
*/
void insert(string word) {
TrieNode* current = root;
for (char ch : word) {
int index = getIndex(ch);
// Create a new node when the current character path is missing.
if (current->children[index] == nullptr) {
current->children[index] = new TrieNode();
}
current = current->children[index];
// Count how many words pass through this node as a prefix.
current->prefixCount++;
}
// Count one more full occurrence of the exact word.
current->wordCount++;
}
/*
Returns how many times the exact word exists in the Trie.
*/
int countWordsEqualTo(string word) {
TrieNode* current = root;
for (char ch : word) {
int index = getIndex(ch);
// If the path breaks, this exact word was never stored.
if (current->children[index] == nullptr) {
return 0;
}
current = current->children[index];
}
return current->wordCount;
}
/*
Returns how many stored words begin with the given prefix.
*/
int countWordsStartingWith(string prefix) {
TrieNode* current = root;
for (char ch : prefix) {
int index = getIndex(ch);
// If the prefix path does not exist, no word can match it.
if (current->children[index] == nullptr) {
return 0;
}
current = current->children[index];
}
return current->prefixCount;
}
/*
Removes one occurrence of a word if it exists in the Trie.
*/
void erase(string word) {
// Erase should change nothing when the word is absent.
if (countWordsEqualTo(word) == 0) {
return;
}
TrieNode* current = root;
for (char ch : word) {
int index = getIndex(ch);
current = current->children[index];
// Reduce the prefix count because one word is being removed.
current->prefixCount--;
}
// Reduce the count of exact word endings by one occurrence.
current->wordCount--;
}
};

Complexity Analysis

Time Complexity: O(L) for insert, countWordsEqualTo, countWordsStartingWith, and erase, where L is the length of the word or prefix.

Space Complexity: O(26 * N) in pointer capacity terms for N nodes in the fixed-array Trie, while one insertion may create up to O(L) new nodes in the worst case.

Interview follow-up Questions

One counter is needed for exact word endings, and the other is needed for prefix frequency. Without separating them, exact-word count and prefix count would get mixed up.

TrieArraysData Structures

Read Similar Blogs

HashingData Structures

Count Distinct Elements

Given an integer array arr[] , return the count of all distinct elements in the array. Input: arr = [10, 20, 20, 10, 30, 10] Output: 3 Explanation: The different values are 10 , 20 , and 30 . 10 an...

76.5k
Arrays

Row wise traversal

Given a matrix with N rows and M columns, print all elements row by row from left to right. Input: mat = [[1, 2, 3], [4, 5, 6]] Output: 1 2 3 4 5 6 Explanation: First row prints as 1 2 3 , then sec...

61.7k
Arrays

Main Diagonal Print

Given a matrix with N rows and M columns, print all elements present on the main diagonal. Input: mat = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] Output: 1 5 9 Explanation: Main diagonal cells are mat[0][0...

67.1k
Comments0