1804. Implement Trie II (Prefix Tree)

跟Trie1 幾乎一樣

就不再說明一次

節點會去記兩個東西, 走過的次數, 跟字串在這邊end的字數

class TrieNode{
    public:
        TrieNode():count(0), end_count(0){};
        unordered_map<char, TrieNode*> children;
        int count;
        int end_count;
};
class Trie {
public:
    Trie() {
        root = new TrieNode();
    }
    
    void insert(string word) {
        TrieNode *node = root;
        node->count++;
        for(auto w: word){
            if(!node->children.count(w)){
                node->children[w] = new TrieNode;
            }
            node = node->children[w];
            node->count++;
        }
        node->end_count++;
    }
    
    int countWordsEqualTo(string word) {
        TrieNode *node = root;
        for(auto w: word){
            if(!node->children.count(w) || node->children[w]->count == 0){
                return 0;
            }
            node = node->children[w];
        }
        return node->end_count;
    }
    
    int countWordsStartingWith(string prefix) {
        TrieNode *node = root;
        for(auto w: prefix){
            if(!node->children.count(w) || node->children[w]->count == 0){
                return 0;
            }
            node = node->children[w];
        }
        return node->count;
    }
    
    void erase(string word) {
        TrieNode *node = root;
        node->count--;
        for(auto w: word){
            if(!node->children.count(w) || node->children[w]->count == 0){
                return;
            }
            node = node->children[w];
            node->count--;
        }
        node->end_count--;
    }
    TrieNode *root;
};
/**
 * 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);
 */