Cod sursa(job #2191854)

Utilizator jeanFMI - Petcu Ion Cristian jean Data 3 aprilie 2018 21:28:50
Problema Trie Scor 40
Compilator cpp Status done
Runda Arhiva educationala Marime 2.3 kb
#include<iostream>
#include<string>
#include<unordered_map>
#include<stack>
using namespace std;

struct TrieNode {
    int count;
    unordered_map<char, TrieNode*> children;
 
    TrieNode() : count(0) {}
};

class Trie {
    private:
        TrieNode* root;
    public:
        Trie();
        void add(string word);
        void remove(string word);
        int search(string word);
        int query(string word);
};

Trie::Trie() {
    this->root = new TrieNode();
}

void Trie::add(string word) {
    TrieNode* node = this->root;
    for (auto ch: word) {
        if (node->children.find(ch) == node->children.end()) {
            node->children[ch] = new TrieNode();
        }
        node = node->children[ch];
    }
    node->count++;
}

void Trie::remove(string word) {
    stack<pair<TrieNode*, char> > st;
    TrieNode* node = this->root;
    for (auto ch: word) {
        if (node->children.find(ch) == node->children.end()) {
            return;
        }
        st.push(make_pair(node, ch));
        node = node->children[ch];
    }
    if (node->count == 0) {
        return;
    }

    node->count--;
    if (node->count == 0) {
        while (!st.empty()) {
            auto pair = st.top();
            st.pop();
            node = pair.first;

            node->children.erase(pair.second);
            if (node->count > 0 || node->children.size() > 0) {
                break;
            }
        }
    }
}

int Trie::search(string word) {
    TrieNode* node = this->root;
    for (auto ch: word) {
        if (node->children.find(ch) == node->children.end()) {
            return 0;
        }
        node = node->children[ch];
    }
    return node->count;
}

int Trie::query(string word) {
    TrieNode* node = this->root;
    int count = 0;
    for (auto ch: word) {
        if (node->children.find(ch) == node->children.end()) {
            return count;
        }
        count++;
        node = node->children[ch];
    }
    return count;
}


int main () {
    freopen("trie.in", "r", stdin);
    freopen("trie.out", "w", stdout);

    Trie T;
    int op;
    string word;
 
    while(cin >> op) {
        cin >> word;
        if (op == 0) T.add(word);
        else if (op == 1) T.remove(word);
        else if (op == 2) cout << T.search(word) << endl;
        else cout << T.query(word) << endl;
    }

    return 0;
}