1804. Implement Trie II (Prefix Tree)

MediumDesignTrieHash TableString

Problem Description

The problem involves implementing a class for a trie, also known as a prefix tree. A trie is a specialized tree used to manage a set of strings in such a way that queries like insertion, search, prefix counting, and deletion can be performed efficiently. The specific operations to implement in this Trie class are:

Trie: Initialize the trie.
insert: Add a string word to the trie.
countWordsEqualTo: Determine how many times a specific string word appears in the trie.
countWordsStartingWith: Count the number of strings in the trie that start with a given prefix.
erase: Remove a string word from the trie.

The goal is to handle these operations as efficiently as possible.

Intuition

Trie Structure:

The trie is constructed using nodes where each node represents a character in a string. Each node might have up to 26 children (for each letter in the English alphabet) and also stores values to keep track of the number of times a word or prefix occurs.

Insertion:

When inserting a word, we start from the root and move through each character of the word. For each character:

Determine the corresponding index (0 to 25 for 'a' to 'z').
If the child node corresponding to that character does not exist, we create a new Trie node there.
Move to the child node and increment the prefix count (pv) as this node will now be part of one more string with a given prefix.
Once all characters are processed, we increment the word count (v) at the last node to indicate the end of a word.

Counting Words Equal:

To count the instances of a specific word, traverse the trie following the characters of the word similar to insertion:

If any character's corresponding node does not exist, it means the word is not present, and we return 0.
If we reach the end of the word and all nodes were present, we return the word count (v) saved at the last node.

Counting Words Starting With:

Counting words with a given prefix also follows a path down the trie based on the prefix's characters:

If a node corresponding to a character in the prefix does not exist, return 0 as no words start with that prefix.
If the prefix is traversed successfully, return the prefix count (pv) from the last node of the prefix.

Erase Operation:

To erase a word:

Traverse the trie character by character, decrementing the prefix count (pv) along the way, as the trie will now contain one less instance of those prefixes.
At the final node, decrement the word count (v) to indicate one less occurrence of the complete word.

This design uses two count variables (v for the count of exact words and pv for the prefix count) at each trie node to support efficient operations.

Learn more about Trie patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece：

Solution Approach

The implementation of the Trie class uses a simple, yet effective design pattern to operate and manage strings.

Node Structure:

Each node in the trie contains an array called children with a length of 26, representing the possible characters (lowercase 'a' to 'z'). In addition, each node has two integer variables: v for tracking the actual words that end at that node, and pv for tracking the number of words or prefixes passing through that node.

Insert Operation (`insert` Method):

The insert function adds a new word to the trie following these steps:

Loop through each character c in the word.
Calculate the index idx that represents the character in the children array using idx = ord(c) - ord('a').
If node.children[idx] is None, create a new Trie node at that index.
Move to the corresponding child node and increment the prefix value count node.pv.
At the end of the loop, increment node.v which represents the count of the word at that node.

Search Helper Function (`search` Method):

The search helper function is used by both the countWordsEqualTo and countWordsStartingWith functions to traverse the trie based on the input word or prefix. It moves down the trie one character at a time. If a character does not have a corresponding child node, the function returns None. If all characters are found in the trie, it returns the final node reached.

Count Words Equal Operation (`countWordsEqualTo` Method):

This function:

Uses the search function to navigate to the node corresponding to the last character of the word.
If the word exists in the trie, it will return the word count v from the last node.
Otherwise, it returns 0.

Count Words Starting With Operation (`countWordsStartingWith` Method):

This function:

Utilizes the search function to find the node corresponding to the last character of the prefix.
If a valid node is found, it returns the prefix value count pv.
Otherwise, it returns 0.

Erase Operation (`erase` Method):

To remove a word from the trie:

Traverse the trie following the characters of the word.
At each step, reduce the prefix value count pv as we're removing a word that passes through those nodes.
Once we reach the final node of the word, decrement the word count v.

By maintaining separate counts for words and prefixes (v and pv), these operations are made efficient, allowing us to add, count, and remove words and prefixes in optimal time related to the length of the input string. The trie thus serves as an effective data structure to implement an autocomplete system or a spellchecker where such operations are frequently needed.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Consider the classic dynamic programming of fibonacci numbers, what is the recurrence relation?

dp[i] = dp[i - 1] + dp[i - 2]dp[i] = dp[i - 1] + 1dp[i] = dp[i - 1] + dp[i - 2] + dp[i - 3]dp[i] = dp[i - 1]

I don’t know

Example Walkthrough

Let's illustrate the solution approach with an example. Suppose we start with an empty trie, and we want to perform the following operations:

Insert the words "apple", "app", and "apex"
Count how many times the word "app" appears in the trie
Count how many words start with the prefix "ap"
Erase the word "apple"

Insert Operation:

Insert "apple":
- For each character in "apple", calculate the corresponding index in the children array (0 for 'a', 15 for 'p', and so on).
- At each character, if the child node does not exist, create one, and move to it.
- While moving to each child node, increment the pv (prefix value) count.
- After the last character 'e' is processed, increment v (word count) to indicate that "apple" ends here.
Insert "app":
- Proceed similarly as "apple". Nodes for 'a' and the first 'p' already exist, so we follow them.
- Increment pv at each step.
- After processing the second 'p', we increment v because "app" is also a word in the trie.
Insert "apex":
- The nodes for 'a' and 'p' exist. Follow them and increment pv.
- 'e' and 'x' nodes don't exist after 'p', so create them and increment their pv.
- Increment v at 'x' to indicate the end of the word "apex".

Count Words Equal:

Count how many times "app" appears:

Use the search function to traverse the trie with the word "app".
If we successfully find 'a', follow through two 'p's, and we find that node.v is 1 at the last 'p'.
Hence, "app" appears once in the trie.

Count Words Starting With:

Count how many words start with "ap":

Using the search function, find the node for 'a', then 'p'.
At 'p', node.pv might be 3 because we inserted "apple", "app", and "apex", all of which start with "ap".

Erase Operation:

Erase "apple":

Traverse the trie following "apple".
Decrement pv at each node corresponding to 'a', 'p', 'p', 'l'.
Reach 'e', decrement v because "apple" should not be counted as a word anymore.
After this operation, pv for 'p' should now be 2 since there are still two words starting with "ap" (even though "apple" was erased).

This small example demonstrates how we efficiently insert, count, and remove words from a trie using the described methods. The trie structure enables us to perform these operations with time complexity proportional to the length of the word or prefix, making it highly efficient for tasks involving a large number of string manipulations.

Solution Implementation

1class Trie:
2    def __init__(self):
3        # Initialize a trie node with children and value counters
4        self.children = [None] * 26  # 26 English lowercase letters
5        self.word_end_count = 0  # To count complete words
6        self.prefix_count = 0  # To count words that have the current node as a prefix
7
8    def insert(self, word: str) -> None:
9        # Insert a word into the trie
10        node = self
11        for char in word:
12            index = ord(char) - ord('a')
13            if node.children[index] is None:
14                node.children[index] = Trie()
15            node = node.children[index]
16            node.prefix_count += 1  # Increment the prefix count for each node
17        node.word_end_count += 1  # Mark the end of the word
18
19    def countWordsEqualTo(self, word: str) -> int:
20        # Return the number of instances of a word in the trie
21        node = self._search(word)
22        return 0 if node is None else node.word_end_count
23
24    def countWordsStartingWith(self, prefix: str) -> int:
25        # Return the number of words in the trie that start with the given prefix
26        node = self._search(prefix)
27        return 0 if node is None else node.prefix_count
28
29    def erase(self, word: str) -> None:
30        # Remove a word from the trie
31        node = self
32        for char in word:
33            index = ord(char) - ord('a')
34            if node.children[index] is None:
35                return  # Word not found
36            node = node.children[index]
37            node.prefix_count -= 1  # Decrement the prefix count for each node
38        node.word_end_count -= 1  # Decrement the word count at the end
39
40    def _search(self, word: str):
41        # Helper method to search for a node representing a word or prefix in the trie
42        node = self
43        for char in word:
44            index = ord(char) - ord('a')
45            if node.children[index] is None:
46                return None
47            node = node.children[index]
48        return node
49
50# Example usage
51# trie = Trie()
52# trie.insert("apple")
53# count = trie.countWordsEqualTo("apple")  # count should be 1
54# count_prefix = trie.countWordsStartingWith("app")  # count_prefix should be 1
55# trie.erase("apple")
56

1class Trie {
2    private Trie[] children; // array representing the children nodes of the current node
3    private int wordCount; // number of times a complete word terminates at this node
4    private int prefixCount; // number of words sharing the prefix that ends at this node
5
6    // Constructor
7    public Trie() {
8        children = new Trie[26]; // Initialize the children array to hold 26 English lowercase letters
9        wordCount = 0;
10        prefixCount = 0;
11    }
12
13    // Method to insert a word into the trie
14    public void insert(String word) {
15        Trie node = this; // Start from the root node
16        for (char c : word.toCharArray()) {
17            int index = c - 'a'; // Convert char to index (0-25)
18            if (node.children[index] == null) {
19                node.children[index] = new Trie(); // Create a new Trie node if it does not exist
20            }
21            node = node.children[index]; // Move to the child node
22            node.prefixCount++; // Increment the prefix count for each character
23        }
24        node.wordCount++; // Once the end of the word is reached, increment the word count
25    }
26
27    // Method to count the number of times a word appears in the trie
28    public int countWordsEqualTo(String word) {
29        Trie node = search(word); // Utilize the search helper function
30        return node == null ? 0 : node.wordCount; // If the word is not found, return 0; otherwise, return the word count
31    }
32
33    // Method to count the number of words that start with a given prefix in the trie
34    public int countWordsStartingWith(String prefix) {
35        Trie node = search(prefix); // Utilize the search helper function
36        return node == null ? 0 : node.prefixCount; // If the prefix is not found, return 0; otherwise, return the prefix count
37    }
38
39    // Method to remove a word from the trie
40    public void erase(String word) {
41        Trie node = this; // Start from the root node
42        for (char c : word.toCharArray()) {
43            int index = c - 'a'; // Convert char to index (0-25)
44            node = node.children[index]; // Move to the child node
45            node.prefixCount--; // Decrement the prefix count for each character
46        }
47        node.wordCount--; // Once the end of the word is reached, decrement the word count
48    }
49
50    // Helper function to traverse the trie based on the given word or prefix
51    private Trie search(String word) {
52        Trie node = this; // Start from the root node
53        for (char c : word.toCharArray()) {
54            int index = c - 'a'; // Convert char to index (0-25)
55            if (node.children[index] == null) {
56                return null; // If at any point the node doesn't have a required child, return null
57            }
58            node = node.children[index]; // Otherwise, move to the child node
59        }
60        return node; // Return the node where the word or prefix ends
61    }
62}
63
64/**
65 * The Trie object usage:
66 * Trie trie = new Trie();
67 * trie.insert(word); // Insert a word into the trie
68 * int equalCount = trie.countWordsEqualTo(word); // Get the number of times a word appears in the trie
69 * int prefixCount = trie.countWordsStartingWith(prefix); // Get the number of words starting with a given prefix
70 * trie.erase(word); // Remove a word from the trie
71 */
72

1#include <vector>
2#include <string>
3
4class Trie {
5public:
6    // Constructor initializes the Trie with the 26 English lowercase letters.
7    Trie()
8        : children(26) // Reserve space for 26 children pointers for 'a' - 'z'.
9        , value(0) // Represents the number of times this ending node is used for inserted words.
10        , prefixValue(0) { // Represents the number of times nodes in this path are traversed, i.e., prefix count.
11    }
12
13    // Insert a word into the Trie.
14    void insert(std::string word) {
15        Trie* node = this;
16        for (char c : word) {
17            c -= 'a'; // Convert char to index (0-25).
18            if (!node->children[c]) {
19                node->children[c] = new Trie(); // Create a new node if it does not exist.
20            }
21            node = node->children[c];
22            ++node->prefixValue; // Increment the prefix count of the current node.
23        }
24        ++node->value; // Increment the word ending count.
25    }
26
27    // Count the exact number of occurrences of a specific word in the Trie.
28    int countWordsEqualTo(std::string word) {
29        Trie* node = search(word);
30        return node ? node->value : 0; // If the word is found, return its count, otherwise return 0.
31    }
32
33    // Count the number of words in the Trie that start with the given prefix.
34    int countWordsStartingWith(std::string prefix) {
35        Trie* node = search(prefix);
36        return node ? node->prefixValue : 0; // If the prefix is found, return its count, otherwise return 0.
37    }
38
39    // Erase a word from the Trie.
40    void erase(std::string word) {
41        Trie* node = this;
42        for (char c : word) {
43            c -= 'a'; // Convert char to index (0-25).
44            node = node->children[c];
45            --node->prefixValue; // Decrement the prefix count of the current node.
46        }
47        --node->value; // Decrement the word ending count.
48    }
49
50private:
51    std::vector<Trie*> children; // Pointers to children nodes of the Trie, representing 'a' to 'z'.
52    int value; // Number of times this node marks the end of a word.
53    int prefixValue; // Number of times nodes in this route are part of a prefix.
54
55    // Helper function to traverse the Trie to the node corresponding to the end of the given string.
56    Trie* search(std::string& word) {
57        Trie* node = this;
58        for (char c : word) {
59            c -= 'a';
60            if (!node->children[c]) {
61                return nullptr; // Return null if the character path doesn't exist.
62            }
63            node = node->children[c];
64        }
65        return node; // Return the final node in the path for the word or prefix.
66    }
67};
68

1// Define the maximum number of children (26 for the English alphabet)
2const ALPHABET_SIZE: number = 26;
3
4// Define a node structure for Trie
5interface TrieNode {
6  children: (TrieNode | null)[];
7  value: number;
8  prefixValue: number;
9}
10
11// Initialize Trie with an empty node
12let root: TrieNode = createNode();
13
14// Create a new Trie node with default values
15function createNode(): TrieNode {
16  return {
17    children: Array(ALPHABET_SIZE).fill(null),
18    value: 0,
19    prefixValue: 0,
20  };
21}
22
23// Insert a word into the Trie
24function insert(word: string): void {
25  let node: TrieNode = root;
26  for (let char of word) {
27    let index = char.charCodeAt(0) - 'a'.charCodeAt(0);
28    if (!node.children[index]) {
29      node.children[index] = createNode();
30    }
31    node = node.children[index]!;
32    node.prefixValue++;
33  }
34  node.value++;
35}
36
37// Search for a word in the Trie and return the end node of the word
38function search(word: string): TrieNode | null {
39  let node: TrieNode = root;
40  for (let char of word) {
41    let index = char.charCodeAt(0) - 'a'.charCodeAt(0);
42    if (!node.children[index]) {
43      return null;
44    }
45    node = node.children[index]!;
46  }
47  return node;
48}
49
50// Count the exact number of occurrences of a specific word in the Trie
51function countWordsEqualTo(word: string): number {
52  let node = search(word);
53  return node ? node.value : 0;
54}
55
56// Count the number of words in the Trie that start with the given prefix
57function countWordsStartingWith(prefix: string): number {
58  let node = search(prefix);
59  return node ? node.prefixValue : 0;
60}
61
62// Erase a word from the Trie
63function erase(word: string): void {
64  let node: TrieNode = root;
65  for (let char of word) {
66    let index = char.charCodeAt(0) - 'a'.charCodeAt(0);
67    node = node.children[index]!;
68    node.prefixValue--;
69  }
70  node.value--;
71}
72

Not Sure What to Study? Take the 2-min Quiz：

Is the following code DFS or BFS?

1void search(Node root) {
2  if (!root) return;
3  visit(root);
4  root.visited = true;
5  for (Node node in root.adjacent) {
6    if (!node.visited) {
7      search(node);
8    }
9  }
10}

Depth First SearchBreadth First Search

I don’t know

Time and Space Complexity

Time Complexity

insert method: This involves iterating over each character of the word, hence the time complexity is O(n) where n is the length of the word being inserted. Inserting each character involves a constant amount of work, indexing the array.
countWordsEqualTo method: This simply calls the search method, which traverses the Trie based on the word length, so the time complexity here is O(n), where n is the length of the word we are checking.
countWordsStartingWith method: Similarly, this calls the search method and then returns the pv value associated with the last node. The time complexity of the search portion is O(m) where m is the length of the prefix. Therefore, the overall time complexity is O(m).
erase method: The erase method also traverses the Trie, decrementing the pv count for each node along the path of the word and finally decrements the v count of the final node. Thus, the time complexity is O(n), with n being the length of the word.
search method: The search method iterates over each character of the word and traverses down the Trie accordingly. This operation takes O(n) time where n is the length of the word or prefix being searched.

Space Complexity

Trie structure: Each Trie node has an array of 26 children, so each node uses O(26) = O(1) space. However, the total space used by the Trie depends on the number of nodes and the number of characters in all inserted words. If there are 'w' words with an average length of 'l', the total space complexity would be O(w * l).
insert, countWordsEqualTo, countWordsStartingWith, erase, and search methods: These methods use constant space as they only use pointers and integer operations while traversing the Trie, so the space complexity of each is O(1).

Note: The space used to store input words themselves outside the Trie is not considered in the space complexity of the Trie operations.

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

Fast Track Your Learning with Our Quick Skills Quiz:

Overview

Basic Data Structures

Basic Algorithms

Overview

Sorted Array

Implicitly Sorted Array

Advanced

Speedrun

Introduction

DFS on Tree

Binary Search Tree

Advanced

Combinatorial Search

Pruning

Additional States

Aggregation and Memoization

Dedup

Additional Practices

Speedrun

Introduction

BFS on Tree

Introduction

Vanilla BFS

Matrix as Graph

Implicit Graph

Directed Graph / Topological Sort

Weighted Graph

Minimum Spanning Tree

Speedrun

Introduction

Same Direction

Opposite Direction

Sliding Window

Prefix Sum

Cycle Finding

Advanced

Speedrun

Introduction

Top K

Moving Best

Multiple Heaps

Introduction

Grid

Dynamic number of subproblems

Interval

Dual-Sequence

Knapsack

Bitmask

Additional Practices

Disjoint Set Union | Union Find

Trie

Data Structure Design

Segment Tree

Interval

Monotonic Stack

Divide and Conquer

Line Sweep

Tree Traversal without Recursion

Greedy

Math

Matrix

Classics

Concepts

Projects

Amazon OA

Microsoft OA

Google OA

Twitter OA

Editorials

1804. Implement Trie II (Prefix Tree)

Problem Description

Intuition

Trie Structure:

Insertion:

Counting Words Equal:

Counting Words Starting With:

Erase Operation:

Solution Approach

Node Structure:

Insert Operation (insert Method):

Insert Operation (`insert` Method):

Search Helper Function (`search` Method):

Count Words Equal Operation (`countWordsEqualTo` Method):

Count Words Starting With Operation (`countWordsStartingWith` Method):

Erase Operation (`erase` Method):