题解 | #字典树的实现#

#include <array>
#include <iostream>
using namespace std;

// 根据左神静态数组的思路写的, 不需要动态内存管理了
// https://www.bilibili.com/video/BV1Yu4y1Q7vR/?spm_id_from=333.999.0.0&vd_source=8ea6dc6c02215b66383cf6ab9791e3b4

class TrieTree
{
public:
    static constexpr int MAXN = 150001;  // 操作数是10^5， 还有字符串长度， 理论上要添加的新节点最多1.5*10^5 + 1 ?? 感觉不像？
    std::array<std::array<int, 26>, MAXN> tree;
    std::array<int, MAXN> pass;
    std::array<int, MAXN> end;
    int cnt = 1; // 默认下标0不用

    TrieTree() {
        clear();
    }

    void clear(/*void*/) 
    {
        cnt = 1;
        // 二维数组， 没想到合适一次性清空API TODO 
        for (auto &item : tree) {
            item.fill(0);
        }
        pass.fill(0);
        end.fill(0);
    }

    void insert(const std::string &s)
    {
        int cur = 1;
        pass[cur]++;
        int path;
        for (char ch : s) {
            path = ch - 'a';
            if (tree[cur][path] == 0) {
                tree[cur][path] = ++cnt;
            }
            cur = tree[cur][path];
            pass[cur]++;
        } 
        end[cur]++; // 字符串为空好像也能添加
    }

    // 找到最后一个字符的的节点
    // 找不到返回0 
    int getTreeEnd(const std::string &s)
    {
        int cur = 1;
        int path;
        for (char ch : s)
        {
            path = ch - 'a';
            if (tree[cur][path] == 0) {
                return 0;
            }
            cur = tree[cur][path];
        }
        return cur;
    }

    int search(const std::string &s)
    {
        int cur = getTreeEnd(s); // 逻辑类似， 合并同类项了
        if (cur == 0) {return 0;}
        return end[cur];
    }

    int prefixNumber(const std::string &s)
    {
        int cur = getTreeEnd(s);
        if (cur == 0) {return 0;}
        return pass[cur];
    }

    // delete是关键词，就避免使用了
    void remove(const std::string &s) 
    {
        if (search(s) == 0) {return;} 
        // 前缀树有对应的字符串才删除
        int cur = 1;
        pass[cur]--; // note : 首元素应该也要减的， 不然前缀树上已经存了多少个字符串的信息的错误的
        // 空字符串也是能添加的， 首元素为0， 不能删除首元素， 特殊处理
        int path;
        for (char ch : s)
        {
            // 下一个节点pass减到0了， 后面不在处理
            // 直接遗弃， 对应的节点内存也不使用了
            path = ch - 'a';
            if (--pass[tree[cur][path]] == 0) {
                tree[cur][path] = 0;
                return; // note : 
            }
            cur = tree[cur][path];
        }
        end[cur]--;
    }

};

int main() {
    int n;
    
    int op;
    std::string str;
    TrieTree trie;
    while (std::cin >> n) {
        trie.clear();
        while (n--) {
            std::cin >> op >> str;
            switch (op) {
                case 1: // 添加
                    trie.insert(str);
                    break;
                case 2: // 删除
                    trie.remove(str);
                    break;
                case 3: // 查询str的单词数量
                    std::cout << (trie.search(str) > 0 ? "YES" : "NO" ) << std::endl;
                    break;
                case 4: // 返回str前缀的单词数量
                    std::cout << trie.prefixNumber(str) << std::endl;
                    break;
            }
        }
    }
}
// 64 位输出请用 printf("%lld")