trietree2.cpp

#include <iostream>
#include <string>
#include <unordered_map>

using namespace std;

class Trie {
public:
    struct TrieNode {
        int pass; // 经过这个节点的次数 (原代码叫 path)
        int end;  // 以这个节点结尾的次数
        unordered_map<int, TrieNode*> nexts;

        TrieNode() {
            pass = 0;
            end = 0;
        }

        // 析构函数：递归释放内存
        ~TrieNode() {
            for (auto& pair : nexts) {
                // 修复：删除非空的子节点
                if (pair.second != nullptr) {
                    delete pair.second;
                }
            }
        }
    };

private:
    TrieNode* root;

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

    ~Trie() {
        if (root != nullptr) {
            delete root;
            root = nullptr;
        }
    }

    void insert(const string& word) {
        TrieNode* node = root;
        node->pass++; // 根节点的 pass 也要增加
        for (int i = 0; i < word.length(); i++) {
            int path = word[i] - 'a'; // 这里用 path 表示路径索引
            if (node->nexts.find(path) == node->nexts.end()) {
                node->nexts[path] = new TrieNode();
            }
            node = node->nexts[path];
            node->pass++;
        }
        node->end++;
    }

    void erase(const string& word) {
        if (countWordEqualTo(word) > 0) { // 确认存在才删除
            TrieNode* node = root;
            node->pass--; // 修复：根节点的 pass 也要减少
            for (int i = 0; i < word.length(); i++) {
                int path = word[i] - 'a';
                TrieNode* next = node->nexts[path];
                
                if (--next->pass == 0) {
                    // 如果下个节点的引用计数为0，直接释放内存并从 map 移除
                    delete next;
                    node->nexts.erase(path);
                    return;
                }
                node = next;
            }
            node->end--;
        }
    }

    // 查询单词出现的次数
    int countWordEqualTo(const string& word) {
        TrieNode* node = root;
        for (int i = 0; i < word.length(); i++) {
            int path = word[i] - 'a';
            if (node->nexts.find(path) == node->nexts.end()) {
                return 0;
            }
            node = node->nexts[path];
        }
        return node->end;
    }

    // 查询以前缀开头的单词数量
    int countWordsStartingWith(const string& pre) {
        TrieNode* node = root;
        for (int i = 0; i < pre.length(); i++) {
            int path = pre[i] - 'a';
            if (node->nexts.find(path) == node->nexts.end()) {
                return 0;
            }
            node = node->nexts[path];
        }
        // 修复：这里返回 pass (经过次数)，而不是 end
        return node->pass;
    }
};

// 测试用的 main 函数
int main() {
    Trie trie;
    trie.insert("apple");
    trie.insert("apple");
    trie.insert("app");
    
    cout << "Count 'apple': " << trie.countWordEqualTo("apple") << endl; // 应输出 2
    cout << "Start with 'app': " << trie.countWordsStartingWith("app") << endl; // 应输出 3
    
    trie.erase("apple");
    cout << "Count 'apple' after erase: " << trie.countWordEqualTo("apple") << endl; // 应输出 1
    
    return 0;
}