Problem Statement

Implement "TRIE” data structure from scratch with the following functions.

  • Trie(): Initialize the object of this “TRIE” data structure.
  • insert(“WORD”): Insert the string “WORD” into this “TRIE” data structure.
  • countWordsEqualTo(“WORD”): Return how many times this “WORD” is present in this “TRIE”.
  • countWordsStartingWith(“PREFIX”): Return how many words are there in this “TRIE” that have the string “PREFIX” as a prefix.
  • erase(“WORD”): Delete one occurrence of the string “WORD” from the “TRIE”.

Examples

Example 1:

Input : ["Trie", "insert", "countWordsEqualTo", "insert", "countWordsStartingWith", "erase", "countWordsStartingWith"]
[ [], ["apple"], ["apple"], ["app"], ["app"], ["apple"], ["app"] ]
Output : [null, null, 1, null, 2, null, 1]

Explanation :
Trie trie = new Trie()
trie.insert("apple")
trie.countWordsEqualTo("apple")  // return 1
trie.insert("app") 
trie.countWordsStartingWith("app") // return 2
trie.erase("apple")
trie.countWordsStartingWith("app")   // return 1

Different Approaches

1️⃣ Trie Data Structure

Approach:

To Insert a Node in the Prefix Trie:

  1. Begin at the root node.
  2. Iterate through each character in the word. For each character:
    1. Verify if a child node for that character exists.
    2. If it doesn't, create a new node for the character and link it to the current node.
    3. Increase the prefix count for each node.
  3. Finally, increase the end count for the last node of the word.

To Count the Number of Words Equal to the Given Word in the Trie:

  1. Begin at the root node.
  2. Iterate through each character in the word. For each character:
    1. Traverse down the Trie by following the characters of the word.
    2. If any character is not found, return 0 (indicating the word is not in the Trie).
  3. Once all characters are found, return the end count of the node corresponding to the last character.

To Count the Number of Words Starting with the Given Word in the Trie:

  1. Begin at the root node.
  2. Iterate through each character in the prefix. For each character:
    1. Move down the Trie by following the characters of the prefix.
    2. If any character is not found, return 0 (indicating no words start with the given prefix).
  3. Once all characters are found, return the prefix count of the node corresponding to the last character.

To Erase the Given Word from the Trie:

  1. Begin at the root node.
  2. Iterate through each character in the word. For each character, move down the Trie following the characters of the word and at each node:
    1. Decrease the prefix count.
    2. It is assumed that the word to be erased is always present in the Trie.
  3. Once the end of the word is reached, decrease the end count.

Code:

#include <iostream>
#include <string>
using namespace std;

class TrieNode {
public:
    TrieNode* children[26];  // Pointers for each letter 'a' to 'z'
    bool isEndOfWord;        // Flag to indicate if a word ends here
    int prefixCount;         // Count of words passing through this node (including those ending here)
    int wordCount;           // Count of words that exactly end at this node

    TrieNode() : isEndOfWord(false), prefixCount(0), wordCount(0) {
        for (int i = 0; i < 26; i++) {
            children[i] = nullptr;
        }
    }
};

class Trie {
private:
    TrieNode* root;  // Root node of the trie

public:
    Trie() {
        root = new TrieNode();
    }

    // Inserts a word into the trie
    void insert(const string& word) {
        TrieNode* current = root;
        for (char ch : word) {
            int index = ch - 'a';
            // Create the node if it doesn't exist
            if (current->children[index] == nullptr) {
                current->children[index] = new TrieNode();
            }
            // Increase prefix count for this node as one more word passes through it
            current->children[index]->prefixCount++;
            current = current->children[index];
        }
        // Mark the end of the word and increase the count of words ending here
        current->isEndOfWord = true;
        current->wordCount++;
    }

    // Returns the number of times the exact word has been inserted into the trie
    int countWordsEqualTo(const string& word) {
        TrieNode* current = root;
        for (char ch : word) {
            int index = ch - 'a';
            if (current->children[index] == nullptr) {
                return 0;
            }
            current = current->children[index];
        }
        return current->wordCount;
    }

    // Returns the number of words in the trie that start with the given prefix
    int countWordsStartingWith(const string& prefix) {
        TrieNode* current = root;
        for (char ch : prefix) {
            int index = ch - 'a';
            if (current->children[index] == nullptr) {
                return 0;
            }
            current = current->children[index];
        }
        return current->prefixCount;
    }

    // Erases one occurrence of the word from the trie (reducing counts appropriately)
    void erase(const string& word) {
        TrieNode* current = root;
        for (char ch : word) {
            int index = ch - 'a';
            if (current->children[index] == nullptr) {
                return;  // Word not found; nothing to erase
            }
            // Decrement the prefix count as this word is being removed
            current->children[index]->prefixCount--;
            current = current->children[index];
        }
        // Decrement the count of words ending at this node
        current->wordCount--;
        if (current->wordCount == 0) {
            current->isEndOfWord = false;
        }
    }
};

int main() {
    Trie trie;
    
    // Inserting words into the trie
    trie.insert("apple");
    trie.insert("apple");
    trie.insert("app");
    
    cout << "Count for word 'apple': " << trie.countWordsEqualTo("apple") << endl;
    cout << "Count for word 'app': " << trie.countWordsEqualTo("app") << endl;
    cout << "Count for prefix 'app': " << trie.countWordsStartingWith("app") << endl;
    
    // Erase one occurrence of "apple"
    trie.erase("apple");
    cout << "Count for word 'apple' after erasing once: " << trie.countWordsEqualTo("apple") << endl;
    cout << "Count for prefix 'app' after erasing once: " << trie.countWordsStartingWith("app") << endl;
    
    return 0;
}

Complexity Analysis:

  • Time Complexity:O(N) where N is the length of the word or prefix being processed.
    • Each method (inserting a word, counting words equal to a given word, counting words starting with a prefix, and erasing a word) involves traversing the Trie for each character of the input word or prefix.
    • Thus, the time complexity is linear relative to the length of the word or prefix being processed.
  • Space Complexity:O(N) where N is the total number of characters across all words inserted into the Trie. The space complexity depends on the number of unique words added to the Trie and the average length of these words.