题解 | #字典树的实现#

#include <iostream>
#include <vector>
#include <string>

using namespace std;

struct TrieNode
{
    int pass;
    int end;
    vector<TrieNode*> nexts;

    TrieNode()
        :pass(0),end(0),nexts(26,nullptr) {};
};

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

    ~Trie()
    {
        freeTrie(root);
    }    

    void insert(const string& word)
    {
        TrieNode* node = root;
        node->pass++;
        for(auto ch : word)
        {
            int path = ch - 'a';
            if(path < 0 || path >= 26) return;
            if(node->nexts[path] == nullptr)
            {
                node->nexts[path] = new TrieNode();
            }
            node = node->nexts[path];
            node->pass++;
        }
        node->end++;
    }

    void erase(const string& word)
    {
        if(search(word) > 0)
        {
            TrieNode* node = root;
            node->pass--;
            for(auto ch : word)
            {
                int path = ch - 'a';
                if(path < 0 || path >= 26) return ;
                if(--node->nexts[path]->pass == 0)
                {
                    delete node->nexts[path];
                    node->nexts[path] = nullptr;
                    return;
                }
                
                node = node->nexts[path];
            }

            node->end--;
        }
    }

    int search(const string& word)
    {
        TrieNode* node = root;
        for(auto ch : word)
        {
            int path = ch - 'a';
            if(path < 0 || path >=26 || node->nexts[path] == nullptr)
            {
                return 0;
            }
            node = node->nexts[path];
        }
        return node->end ;
    }

    int prefixNumber(string pre)
    {
        TrieNode* node = root;
        for(auto ch : pre)
        {
            int path = ch - 'a';
            if(path < 0 || path >= 26 || node->nexts[path] == nullptr)
            {
                return 0;
            }
            node = node->nexts[path];
        }
        return node->pass;
    }

private:
    void freeTrie(TrieNode* node)
    {
        if(node == nullptr)
        {
            return;
        }

        for(TrieNode* next : node->nexts)
        {
            freeTrie(next);
        }

        delete node;
    }

private:
    TrieNode* root;
};

int main() {
    int op = 0;
    string str;
    int n = 0;
    while(cin >> n)
    {
        Trie trie;
        while(n--)
        {
            cin>>op>>str;
            switch(op)
            {
                case 1:
                    trie.insert(str);
                    break;
                case 2:
                    trie.erase(str);
                    break;
                case 3:
                    cout << (trie.search(str) > 0 ? "YES" : "NO")<<endl;
                    break;
                case 4:
                    cout << trie.prefixNumber(str) <<endl;
                    break;
            }
        }
    }
}
// 64 位输出请用 printf("%lld")