---
description: Trie
---

# LeetCode 1804. Implement Trie II (Prefix Tree)

A [**trie**](https://en.wikipedia.org/wiki/Trie) (pronounced as "try") or **prefix tree** is a tree data structure used to efficiently store and retrieve keys in a dataset of strings. There are various applications of this data structure, such as autocomplete and spellchecker.

Implement the Trie class:

* `Trie()` Initializes the trie object.
* `void insert(String word)` Inserts the string `word`into the trie.
* `int countWordsEqualTo(String word)` Returns the number of instances of the string `word` in the trie.
* `int countWordsStartingWith(String prefix)`Returns the number of strings in the trie that have the string `prefix` as a prefix.
* `void erase(String word)` Erases the string `word`from the trie.

**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");               // Inserts "apple".
trie.insert("apple");               // Inserts another "apple".
trie.countWordsEqualTo("apple");    // There are two instances of "apple" so return 2.
trie.countWordsStartingWith("app"); // "app" is a prefix of "apple" so return 2.
trie.erase("apple");                // Erases one "apple".
trie.countWordsEqualTo("apple");    // Now there is only one instance of "apple" so return 1.
trie.countWordsStartingWith("app"); // return 1
trie.erase("apple");                // Erases "apple". Now the trie is empty.
trie.countWordsStartingWith("app"); // return 0
```

**Constraints:**

* `1 <= word.length, prefix.length <= 2000`
* `word` and `prefix` consist only of lowercase English letters.
* At most `3 * 10^4` calls **in total** will be made to `insert`, `countWordsEqualTo`, `countWordsStartingWith`, and `erase`.
* It is guaranteed that for any function call to `erase`, the string `word` will exist in the trie.

## Solution

[English Version in Youtube](https://youtu.be/_3ReIHJeTT4)

[中文版解答Youtube Link](https://youtu.be/PAvVIROaw-Q)

[中文版解答Bilibili Link](https://www.bilibili.com/video/BV1jv41187qN/)

```
class Trie {
    
    struct TrieNode {
        TrieNode* nodes[26];
        int word_cnt;
        int prefix_cnt;
        // Initialize your data structure here.
        TrieNode() {
            memset(nodes, 0, sizeof(nodes));
            word_cnt = 0;
            prefix_cnt = 0;
        }
    };
    
    TrieNode* root;
    
public:
    Trie() {
        root = new TrieNode();
    }
    
    void insert(string word) {
        TrieNode* tmp = root;
        for(char ch : word) {
            int index = ch - 'a';
            if(tmp->nodes[index] == nullptr) {
                tmp->nodes[index] = new TrieNode();
            }
            tmp = tmp->nodes[index];
            tmp->prefix_cnt++;
        }
        tmp->word_cnt++;
    }
    
    int countWordsEqualTo(string word) {
        TrieNode* tmp = root;
        for(char ch : word) {
            int index = ch - 'a';
            if(tmp->nodes[index] == nullptr) {
                return 0;
            }
            tmp = tmp->nodes[index];
        }
        return tmp->word_cnt;
    }
    
    int countWordsStartingWith(string prefix) {
        TrieNode* tmp = root;
        for(char ch : prefix) {
            int index = ch - 'a';
            if(tmp->nodes[index] == nullptr) {
                return 0;
            }
            tmp = tmp->nodes[index];
        }
        return tmp->prefix_cnt;
    }
    
    void erase(string word) {
        TrieNode* tmp = root;
        TrieNode* to_be_deleted = nullptr;
        for(char ch : word) {
            int index = ch - 'a';
            TrieNode* parent = tmp;
            tmp = tmp->nodes[index];
            tmp->prefix_cnt--;
            
            if (to_be_deleted != nullptr) {
                delete to_be_deleted;
            }
            if (tmp->prefix_cnt == 0) {
                if (to_be_deleted == nullptr) {
                    parent->nodes[index] = nullptr;
                }
                to_be_deleted = tmp;
            }
        }
        tmp->word_cnt--;
        if (to_be_deleted != nullptr) {
            delete to_be_deleted;
        }
    }
};

/**
 * Your Trie object will be instantiated and called as such:
 * Trie* obj = new Trie();
 * obj->insert(word);
 * int param_2 = obj->countWordsEqualTo(word);
 * int param_3 = obj->countWordsStartingWith(prefix);
 * obj->erase(word);
 */
```
