DSA
Implement Trie - II
Trie problem — solution with code and analysis.
Practice Link
Ninja has to implement a data structure ”TRIE” from scratch. Ninja has to complete some functions.
-
Trie(): Ninja has to initialize the object of this “TRIE” data structure.
-
insert(“WORD”): Ninja has to insert the string “WORD” into this “TRIE” data structure.
-
countWordsEqualTo(“WORD”): Ninja has to return how many times this “WORD” is present in this “TRIE”.
-
countWordsStartingWith(“PREFIX”): Ninjas have to return how many words are there in this “TRIE” that have the string “PREFIX” as a prefix.
-
erase(“WORD”): Ninja has to delete one occurrence of the string “WORD” from the “TRIE”. Note:
-
If erase(“WORD”) function is called then it is guaranteed that the “WORD” is present in the “TRIE”.
-
If you are going to use variables with dynamic memory allocation then you need to release the memory associated with them at the end of your solution. Can you help Ninja implement the "TRIE" data structure?
-
Design#
This extends the basic Trie by adding two counters to each node: wordCount (how many complete words end at this node) and prefixCount (how many words pass through this node). On insert, increment prefixCount at every node along the path and wordCount only at the terminal node. On erase, decrement both counters along the path without deleting nodes (simplifying implementation at the cost of not freeing memory). Queries simply read the appropriate counter at the end of the traversal — all operations remain O(m) where m is word/prefix length.
Implementation#
#include <bits/stdc++.h>
class TrieNode{
public:
TrieNode* children[26];
int wordCount;
int prefixCount;
TrieNode(){
wordCount = 0;
prefixCount = 0;
for(int i=0;i<26;i++)
{
children[i] = NULL;
}
}
};
class Trie{
public:
TrieNode* root;
Trie(){
root = new TrieNode();
}
void insert(string &word){
TrieNode* curr = root;
for(int i=0;i<word.length();i++)
{
int idx = word[i] - 'a';
if(!curr->children[idx])
curr->children[idx] = new TrieNode();
curr =curr->children[idx];
curr->prefixCount++;
}
curr->wordCount++;
}
int countWordsEqualTo(string &word){
TrieNode* curr = root;
for(int i=0;i<word.length();i++)
{
int idx = word[i] - 'a';
if(!curr->children[idx])
return 0;
curr =curr->children[idx];
}
return curr->wordCount;
}
int countWordsStartingWith(string &word){
TrieNode* curr = root;
for(int i=0;i<word.length();i++)
{
int idx = word[i] - 'a';
if(!curr->children[idx])
return 0;
curr =curr->children[idx];
}
return curr->prefixCount;
}
void erase(string &word){
TrieNode* curr = root;
for(int i=0;i<word.length();i++)
{
int idx = word[i] - 'a';
if (curr->children[idx]) {
curr = curr->children[idx];
curr->prefixCount--;
}else
return;
}
curr->wordCount--;
}
};
Complexities#
| Operation | Time Complexity | Space Complexity |
|---|---|---|
| Insert a word | O(N) | O(N*M) |
| Count words equal to a word | O(N) | O(1) |
| Count words starting with a prefix | O(N) | O(1) |
| Erase a word | O(N) | O(1) (without cleanup), O(N) (with cleanup) |