实现 Trie （前缀树） II

Leetcode题目链接

描述

详见中文题目链接。

解答

struct TrieNode {
    TrieNode* nexts[26];
    int passing = 0;
    int ending = 0;

    ~TrieNode() {
        for (auto& node : nexts) {
            delete node;
        }
    }
};

class Trie {
public:
    TrieNode* root;

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

    ~Trie() {
        delete root;
    }

    void insert(string word) {
        TrieNode* node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            if (!node->nexts[index]) {
                node->nexts[index] = new TrieNode();
            }

            node = node->nexts[index];
            node->passing++;
        }

        node->ending++;
    }

    int countWordsEqualTo(string word) {
        TrieNode* node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            if (!node->nexts[index]) {
                return 0;
            }

            node = node->nexts[index];
        }

        return node->ending;
    }

    int countWordsStartingWith(string prefix) {
        TrieNode* node = root;

        for (const char& ch : prefix) {
            int index = ch - 'a';
            if (!node->nexts[index]) {
                return 0;
            }

            node = node->nexts[index];
        }

        return node->passing;
    }

    void erase(string word) {
        TrieNode *node = root;

        for (const char& ch : word) {
            int index = ch - 'a';
            node = node->nexts[index];
            node->passing--;
        }

        node->ending--;
    }
};

/**
 * 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);
 */