1804. Implement Trie II (Prefix Tree) - Leetcode Solution

Code Implementation

class TrieNode:
    def __init__(self):
        self.children = {}
        self.count = 0
        self.end = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = TrieNode()
            node = node.children[ch]
            node.count += 1
        node.end += 1

    def countWordsEqualTo(self, word: str) -> int:
        node = self.root
        for ch in word:
            if ch not in node.children:
                return 0
            node = node.children[ch]
        return node.end

    def countWordsStartingWith(self, prefix: str) -> int:
        node = self.root
        for ch in prefix:
            if ch not in node.children:
                return 0
            node = node.children[ch]
        return node.count

    def erase(self, word: str) -> None:
        node = self.root
        stack = []
        for ch in word:
            if ch not in node.children:
                return  # word not found
            stack.append((node, ch))
            node = node.children[ch]
        node.end -= 1
        # Decrement count along the path
        for parent, ch in stack:
            parent.children[ch].count -= 1

class TrieNode {
public:
    TrieNode* children[26];
    int count;
    int end;
    TrieNode() {
        for (int i = 0; i < 26; ++i) children[i] = nullptr;
        count = 0;
        end = 0;
    }
};

class Trie {
    TrieNode* root;
public:
    Trie() {
        root = new TrieNode();
    }
    
    void insert(string word) {
        TrieNode* node = root;
        for (char ch : word) {
            int idx = ch - 'a';
            if (!node->children[idx])
                node->children[idx] = new TrieNode();
            node = node->children[idx];
            node->count++;
        }
        node->end++;
    }
    
    int countWordsEqualTo(string word) {
        TrieNode* node = root;
        for (char ch : word) {
            int idx = ch - 'a';
            if (!node->children[idx]) return 0;
            node = node->children[idx];
        }
        return node->end;
    }
    
    int countWordsStartingWith(string prefix) {
        TrieNode* node = root;
        for (char ch : prefix) {
            int idx = ch - 'a';
            if (!node->children[idx]) return 0;
            node = node->children[idx];
        }
        return node->count;
    }
    
    void erase(string word) {
        TrieNode* node = root;
        vector<TrieNode*> nodes;
        for (char ch : word) {
            int idx = ch - 'a';
            if (!node->children[idx]) return;
            node = node->children[idx];
            nodes.push_back(node);
        }
        node->end--;
        for (TrieNode* n : nodes) {
            n->count--;
        }
    }
};

class TrieNode {
    TrieNode[] children = new TrieNode[26];
    int count = 0;
    int end = 0;
}

class Trie {
    TrieNode root;

    public Trie() {
        root = new TrieNode();
    }
    
    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            int idx = ch - 'a';
            if (node.children[idx] == null)
                node.children[idx] = new TrieNode();
            node = node.children[idx];
            node.count++;
        }
        node.end++;
    }
    
    public int countWordsEqualTo(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            int idx = ch - 'a';
            if (node.children[idx] == null)
                return 0;
            node = node.children[idx];
        }
        return node.end;
    }
    
    public int countWordsStartingWith(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            int idx = ch - 'a';
            if (node.children[idx] == null)
                return 0;
            node = node.children[idx];
        }
        return node.count;
    }
    
    public void erase(String word) {
        TrieNode node = root;
        List<TrieNode> nodes = new ArrayList<>();
        for (char ch : word.toCharArray()) {
            int idx = ch - 'a';
            if (node.children[idx] == null)
                return;
            node = node.children[idx];
            nodes.add(node);
        }
        node.end--;
        for (TrieNode n : nodes) {
            n.count--;
        }
    }
}

class TrieNode {
    constructor() {
        this.children = {};
        this.count = 0;
        this.end = 0;
    }
}

class Trie {
    constructor() {
        this.root = new TrieNode();
    }
    
    insert(word) {
        let node = this.root;
        for (const ch of word) {
            if (!(ch in node.children)) {
                node.children[ch] = new TrieNode();
            }
            node = node.children[ch];
            node.count += 1;
        }
        node.end += 1;
    }
    
    countWordsEqualTo(word) {
        let node = this.root;
        for (const ch of word) {
            if (!(ch in node.children)) return 0;
            node = node.children[ch];
        }
        return node.end;
    }
    
    countWordsStartingWith(prefix) {
        let node = this.root;
        for (const ch of prefix) {
            if (!(ch in node.children)) return 0;
            node = node.children[ch];
        }
        return node.count;
    }
    
    erase(word) {
        let node = this.root;
        let stack = [];
        for (const ch of word) {
            if (!(ch in node.children)) return;
            stack.push([node, ch]);
            node = node.children[ch];
        }
        node.end -= 1;
        for (const [parent, ch] of stack) {
            parent.children[ch].count -= 1;
        }
    }
}

Problem Description

The task is to implement a Trie (Prefix Tree) data structure that supports the following operations efficiently:

insert(word): Inserts the string word into the trie.
countWordsEqualTo(word): Returns the number of times the exact string word has been inserted into the trie.
countWordsStartingWith(prefix): Returns the number of words in the trie that start with the string prefix.
erase(word): Removes one occurrence of the string word from the trie. If the word was inserted multiple times, only one occurrence is removed.

Each operation must be efficient, and the trie must be able to handle duplicate words and prefixes. All words consist of lowercase English letters.

Thought Process

To solve this problem, we need to choose a data structure that allows us to efficiently insert, search, and erase words, as well as count words with a given prefix. A trie (prefix tree) is ideal for this because it stores words by branching at each character, allowing efficient prefix operations.

A naive approach would be to store every word in a list and iterate through the list for each operation. However, this would be too slow for large inputs, especially for prefix counting. Instead, a trie enables us to share common prefixes and perform operations in time proportional to the length of the word or prefix.

Additionally, to support counting duplicates and erasing, we need to track how many times each word and prefix appears in the trie. This requires augmenting each trie node with counters.

Solution Approach

We design a trie with nodes that store:

A mapping from character to child node (for branching).
A count field: the number of words passing through this node (for prefix counting).
An end field: the number of words ending exactly at this node (for exact word counting).

The operations are implemented as follows:

Insert: For each character in the word, traverse or create the corresponding child node, incrementing the count at each node. At the end node, increment end.
countWordsEqualTo: Traverse the trie by each character in the word. If any character is missing, return 0. Otherwise, return the end at the last node.
countWordsStartingWith: Traverse the trie by each character in the prefix. If any character is missing, return 0. Otherwise, return the count at the last node.
Erase: Traverse the trie by each character in the word, tracking the path. At the end node, decrement end. Then, for each node along the path, decrement count.

This design ensures that all operations run in O(L) time, where L is the length of the word or prefix.

Example Walkthrough

Let's walk through an example with the following operations:

insert("apple")
insert("apple")
insert("app")
countWordsEqualTo("apple") → returns 2
countWordsStartingWith("app") → returns 3
erase("apple")
countWordsEqualTo("apple") → returns 1
countWordsStartingWith("app") → returns 2

Step-by-step:

After inserting "apple" twice and "app" once, the trie has "apple" (count 2) and "app" (count 1), and all relevant prefix counts are updated.
countWordsEqualTo("apple") returns 2 because "apple" was inserted twice.
countWordsStartingWith("app") returns 3 because both "apple" (twice) and "app" (once) start with "app".
After erase("apple"), only one occurrence of "apple" remains, and all prefix counts are decremented accordingly.
countWordsEqualTo("apple") now returns 1, and countWordsStartingWith("app") returns 2.

This demonstrates how the trie efficiently tracks both word and prefix counts, even with duplicates and deletions.

Time and Space Complexity

Brute-force approach:

Time: Each operation would require iterating through all stored words, leading to O(N * L) per operation, where N is the number of words and L is the average word length.
Space: O(N * L) for storing all words individually.

Optimized Trie approach:

Time: Each operation (insert, count, erase) takes O(L), where L is the length of the word or prefix, since we only traverse the trie along the word's characters.
Space: In the worst case (all unique words), space is O(N * L) for the trie nodes. However, shared prefixes save space compared to brute-force.

Summary

The solution leverages a trie data structure to efficiently store, count, and erase words and prefixes. By augmenting trie nodes with counters for words passing through and ending at each node, we achieve O(L) time complexity for all operations. This approach is both elegant and practical, especially for problems involving dynamic prefix queries and duplicate word management.

AlgoMap.io

Want Help Cracking FAANG?