Implement Trie II (Prefix Tree)

Updated on 03 June, 2025
Implement Trie II (Prefix Tree) header image

Problem Statement

You need to implement a trie (prefix tree), a specialized tree data structure used to efficiently store and query strings. This is useful in applications like autocomplete, spell checkers, and string dictionaries.

Your Trie class should support the following methods:

  • Trie(): Initializes the trie object.
  • void insert(String word): Inserts the string word into the trie. If the word already exists, its instance count increases.
  • int countWordsEqualTo(String word): Returns how many times the exact word exists in the trie.
  • int countWordsStartingWith(String prefix): Returns how many words in the trie start with the given prefix.
  • void erase(String word): Removes one occurrence of word from the trie. If the word exists multiple times, only one instance is erased.

Examples

Example 1

Input:

["Trie", "insert", "insert", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsEqualTo", "countWordsStartingWith", "erase", "countWordsStartingWith"]
[[], ["apple"], ["apple"], ["apple"], ["app"], ["apple"], ["apple"], ["app"], ["apple"], ["app"]]

Output:

[null, null, null, 2, 2, null, 1, 1, null, 0]

Explanation:

Trie trie = new Trie();
trie.insert("apple");                      // Adds "apple"
trie.insert("apple");                      // Adds another "apple"
trie.countWordsEqualTo("apple");           // Returns 2
trie.countWordsStartingWith("app");        // Returns 2
trie.erase("apple");                       // Removes one "apple"
trie.countWordsEqualTo("apple");           // Returns 1
trie.countWordsStartingWith("app");        // Returns 1
trie.erase("apple");                       // Removes the last "apple"
trie.countWordsStartingWith("app");        // Returns 0

Constraints

  • 1 <= word.length, prefix.length <= 2000
  • word and prefix consist only of lowercase English letters.
  • At most 3 * 10^4 calls will be made to all operations combined.
  • It is guaranteed that erase(word) is only called when word exists in the trie.

Approach and Intuition

The trie data structure supports all required operations efficiently by leveraging character-wise traversal and storage:

  1. Initialization:

    • The root node starts empty. Each node stores:

      • A map/array of child nodes (one for each character)
      • A count of words ending at that node
      • A count of words passing through that node (for prefix counting)
  2. Insertion:

    • Traverse the trie character by character.

    • For each character:

      • Create a node if it doesn't exist.
      • Increment prefix count.
    • At the final node, increment the end count.

  3. Counting Exact Matches:

    • Traverse through characters in word.
    • If traversal completes, return the node’s end count; otherwise, return 0.
  4. Counting Prefix Matches:

    • Traverse through characters in prefix.
    • Return the prefix count of the last node if traversal succeeds; otherwise, return 0.
  5. Erase:

    • Traverse down the trie as in insertion.
    • Decrement the prefix count along the path.
    • Decrement the end count at the final node.

This design ensures all operations run in O(L) time, where L is the length of the word or prefix, making it efficient even with tens of thousands of operations.

Solutions

  • C++
  • Java
  • JavaScript
  • Python
cpp
class TrieNode {
public:
    TrieNode* children[26];
    int terminalCount = 0;
    int prefixCount = 0;
};

class Trie {
public:
    TrieNode* root;

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

    void addWord(string word) {
        TrieNode* current = root;
        for (char& c : word) {
            int idx = c - 'a';
            if (!current->children[idx]) {
                current->children[idx] = new TrieNode();
            }
            current = current->children[idx];
            current->prefixCount++;
        }
        current->terminalCount++;
    }

    int countExactWords(string word) {
        TrieNode* current = root;
        for (char& c : word) {
            int idx = c - 'a';
            if (!current->children[idx]) {
                return 0;
            }
            current = current->children[idx];
        }
        return current->terminalCount;
    }

    int countPrefixWords(string prefix) {
        TrieNode* current = root;
        for (char& c : prefix) {
            int idx = c - 'a';
            if (!current->children[idx]) {
                return 0;
            }
            current = current->children[idx];
        }
        return current->prefixCount;
    }

    void removeWord(string word) {
        TrieNode* current = root;
        for (char& c : word) {
            int idx = c - 'a';
            current = current->children[idx];
            current->prefixCount--;
        }
        current->terminalCount--;
    }
};

This implementation in C++ provides a detailed example of how to build and manage a Trie (prefix tree), specifically designed for handling operations like adding words, counting exact occurrences of words, counting words that start with a specific prefix, and removing words.

  • The TrieNode class represents each node in the Trie. Each node includes:

    • An array children[26] for storing pointers to child nodes, symbolizing the 26 lowercase letters of the English alphabet.
    • terminalCount, an integer counting how many times a word terminating at this node has been added to the Trie.
    • prefixCount, an integer counting how many words passing through this node have been added to the Trie.
  • The Trie class facilitates various operations:

    1. The constructor initializes the Trie with a root node.
    2. addWord(string word) inserts a word into the Trie. This function iterates through each character of the word, navigating or creating the necessary child nodes, and increments the prefixCount at each node along the path. At the terminal node of the word, it increments the terminalCount.
    3. countExactWords(string word) returns the number of times a complete word was added to the Trie. It traverses the Trie according to the word's characters, returning 0 if a path breaks, or the terminalCount at the final node if it reaches the end.
    4. countPrefixWords(string prefix) calculates how many words in the Trie start with a specific prefix by similar traversal as countExactWords, but instead of terminalCount, it returns the prefixCount at the last node.
    5. removeWord(string word) decrements the counts related to a previously added word. As it navigates through the Trie following the word's characters, it decrements the prefixCount for each node and terminalCount at the end node.

This Trie implementation is efficient for dictionary-like data manipulations where you often need to check for word existence, frequency, and prefix-related queries.

java
class PrefixTreeNode {
    PrefixTreeNode[] children = new PrefixTreeNode[26];
    int endOfWordCount = 0;
    int prefixCount = 0;
}

class PrefixTree {
    PrefixTreeNode root;

    public PrefixTree() {
        root = new PrefixTreeNode();
    }

    public void insert(String word) {
        PrefixTreeNode node = root;
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                node.children[index] = new PrefixTreeNode();
            }
            node = node.children[index];
            node.prefixCount++;
        }
        node.endOfWordCount++;
    }

    public int countWordsEqualTo(String word) {
        PrefixTreeNode node = root;
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                return 0;
            }
            node = node.children[index];
        }
        return node.endOfWordCount;
    }

    public int countWordsStartingWith(String prefix) {
        PrefixTreeNode node = root;
        for (char c : prefix.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                return 0;
            }
            node = node.children[index];
        }
        return node.prefixCount;
    }

    public void erase(String word) {
        PrefixTreeNode node = root;
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            node = node.children[index];
            node.prefixCount--;
        }
        node.endOfWordCount--;
    }
}

Implement a Trie II (Prefix Tree) in Java to efficiently manage a set of strings and allow for quick retrieval of certain properties of strings, such as the count of words that match a specific prefix or whether a word is present in the trie.

  • Implement the PrefixTreeNode class to represent each node in the trie:

    • Maintain a children array where each element represents a letter of the alphabet.
    • Use endOfWordCount to track how many times a word that ends at this node has been inserted.
    • Use prefixCount to track how many words passing through this node contain the prefix up to this node.
  • In the PrefixTree class, the root instance of PrefixTreeNode acts as the starting point of the Trie:

    • insert: For each character in the word, calculate its position in the alphabet, and navigate to the corresponding child. If the child does not exist, create a new node. Increment prefixCount and, for the last character, increment endOfWordCount.
    • countWordsEqualTo: Traverse the trie using characters of the specified word. If a mismatch is found or the word does not exactly end at a node, return 0. Otherwise, return endOfWordCount for the final node reached.
    • countWordsStartingWith: Similar to countWordsEqualTo, but return prefixCount for the last node of the prefix, indicating the total number of words containing the prefix.
    • erase: Traverse the trie similarly to insert, but decrement the prefixCount for each node along the path, and decrement endOfWordCount at the node where the word terminates.

Implement these methods to ensure efficient word management in the trie, enhancing the performance for adding, removing, and searching words based on prefixes or exact matches. This structure is particularly effective for applications like autocomplete systems, spell-checkers, and other text-processing utilities.

js
class PrefixTreeNode {
    constructor() {
        this.children = new Array(26).fill(null);
        this.endOfWordCount = 0;
        this.prefixCount = 0;
    }
}

var TrieStructure = function() {
    this.root = new PrefixTreeNode();
};

TrieStructure.prototype.addWord = function(word) {
    let currentNode = this.root;
    for (let character of word) {
        const index = character.charCodeAt(0) - 'a'.charCodeAt(0);
        if (!currentNode.children[index]) {
            currentNode.children[index] = new PrefixTreeNode();
        }
        currentNode = currentNode.children[index];
        currentNode.prefixCount++;
    }
    currentNode.endOfWordCount++;
};

TrieStructure.prototype.searchExact = function(word) {
    let currentNode = this.root;
    for (let character of word) {
        const index = character.charCodeAt(0) - 'a'.charCodeAt(0);
        if (!currentNode.children[index]) {
            return 0;
        }
        currentNode = currentNode.children[index];
    }
    return currentNode.endOfWordCount;
};

TrieStructure.prototype.searchPrefix = function(prefix) {
    let currentNode = this.root;
    for (let character of prefix) {
        const index = character.charCodeAt(0) - 'a'.charCodeAt(0);
        if (!currentNode.children[index]) {
            return 0;
        }
        currentNode = currentNode.children[index];
    }
    return currentNode.prefixCount;
};

TrieStructure.prototype.removeWord = function(word) {
    let currentNode = this.root;
    for (let character of word) {
        const index = character.charCodeAt(0) - 'a'.charCodeAt(0);
        currentNode = currentNode.children[index];
        currentNode.prefixCount--;
    }
    currentNode.endOfWordCount--;
};

Implement a Trie (Prefix Tree) in JavaScript to efficiently store and search words by both exact matches and prefixes. The provided solution defines a data structure capable of:

  • Adding words to the Trie.
  • Searching for an exact word in the Trie.
  • Searching for any words in the Trie that have a specific prefix.
  • Removing words from the Trie.

Here's a breakdown of the functionality provided in the code:

  • PrefixTreeNode Class: Represents each node in the Trie with:
    • children: An array of 26 elements initialized to null, representing potential child nodes for each letter of the alphabet.
    • endOfWordCount: An integer to count how many times a word ending at this node has been inserted into the Trie.
    • prefixCount: An integer to count the number of words that use the node as part of their prefix.
  • TrieStructure Class: Represents the entire Trie with:
    • root: The root node of the Trie, initialized as a new instance of PrefixTreeNode.

Methods include:

  1. addWord(word): Inserts a word into the Trie. It increments the prefixCount at each node it visits and sets endOfWordCount at the word's final node.
  2. searchExact(word): Returns the count of how many times a word is found in the Trie exactly as is.
  3. searchPrefix(prefix): Returns the count of how many words start with the given prefix.
  4. removeWord(word): Decrements the prefixCount of each node involved in the word and decrements the endOfWordCount at the word's end node, effectively removing the word if its count reaches zero.

Use this Trie data structure for efficient prefix querying and exact word matching in scenarios like autocomplete features or spell-checkers in JavaScript-based applications.

python
class PrefixNode:
    def __init__(self):
        self.children = [None] * 26
        self.end_count = 0
        self.prefix_count = 0

class PrefixTree:
    def __init__(self):
        self.root = PrefixNode()

    def add_word(self, word: str) -> None:
        node = self.root
        for char in word:
            index = ord(char) - ord('a')
            if not node.children[index]:
                node.children[index] = PrefixNode()
            node = node.children[index]
            node.prefix_count += 1
        node.end_count += 1

    def search_exact(self, word: str) -> int:
        node = self.root
        for char in word:
            index = ord(char) - ord('a')
            if not node.children[index]:
                return 0
            node = node.children[index]
        return node.end_count

    def search_prefix(self, prefix: str) -> int:
        node = self.root
        for char in prefix:
            index = ord(char) - ord('a')
            if not node.children[index]:
                return 0
            node = node.children[index]
        return node.prefix_count

    def remove_word(self, word: str) -> None:
        node = self.root
        for char in word:
            index = ord(char) - ord('a')
            node = node.children[index]
            node.prefix_count -= 1
        node.end_count -= 1

The provided Python code defines a class-based implementation for a specialized data structure known as the Trie (or Prefix Tree). The example specifically covers a PrefixTree class and its operations: adding words, searching for exact words, searching words by prefix, and removing words. This solution effectively manages prefixes and counts the occurrences of each word, which is helpful for tasks such as autocomplete and spell checking.

  • The PrefixNode class initializes with:

    • A list of children nodes to store 26 possible alphabet characters.
    • end_count to track how many times a word ends at this node.
    • prefix_count to mark how many words pass through this node.
  • The PrefixTree class operates with:

    • add_word which increments the prefix count for each character in a word and sets the end count once the entire word is processed.
    • search_exact returns the count of how many times the exact word appears in the trie.
    • search_prefix returns the number of words that start with a given prefix.
    • remove_word to decrement the count of the word and its prefixes from the trie.

For each method:

  1. Nodes navigate through the trie based on the character's position in the alphabet using the formula index = ord(char) - ord('a').
  2. Nodes are dynamically created if they do not exist when trying to store a new character.

This setup ensures efficient processing and retrieval of strings based upon prefixes, widely utilized in systems processing large word databases where quick search operations are critical.

Comments

No comments yet.