Cod sursa(job #2237318)

Utilizator alittlezzCazaciuc Valentin alittlezz Data 1 septembrie 2018 16:28:33
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 3.77 kb
#include <bits/stdc++.h>

using namespace std;

typedef std::unordered_map <char, short> MAP;

class TrieNode{
    public:
        TrieNode();
        TrieNode(short _alphabet_size){
            alphabet_size = _alphabet_size;
            children = new TrieNode*[alphabet_size];
            for(int i = 0;i < alphabet_size;i++){
                children[i] = NULL;
            }
        }
        ~TrieNode(){
            words_touching = 0;
            words_ending = 0;
            delete children;
        }

        int words_touching = 0;
        int words_ending = 0;
        TrieNode **children;
    private:
        short alphabet_size;
};

class Trie{
    public:
        Trie(short alphabet_size, MAP _character_mapping);
        ~Trie();

        void add_word(std::string &word);
        void remove_word(std::string &word);
        void remove_word_recursive(TrieNode *&current_node, std::string &word, int position);
        int count_word(std::string &word);
        int longest_prefix(std::string &word);
    private:
        short alphabet_size;
        MAP character_mapping;
        TrieNode *root;
};

Trie::Trie(short _alphabet_size, MAP _character_mapping){
    alphabet_size = _alphabet_size;
    character_mapping = _character_mapping;
    root = new TrieNode(alphabet_size);
}

Trie::~Trie(){

}

void Trie::add_word(std::string &word){
    TrieNode *current_node = root;
    for(int i = 0;i < word.size();i++){
        TrieNode** child = &current_node->children[character_mapping[word[i]]];
        if((*child) == NULL){
            *child = new TrieNode(alphabet_size);
        }
        (*child)->words_touching++;
        if(i == word.size() - 1){
            (*child)->words_ending++;
        }
        current_node = *child;
    }
}

void Trie::remove_word(std::string &word){
    remove_word_recursive(root, word, 0);
}

void Trie::remove_word_recursive(TrieNode *&current_node, std::string &word, int position){
    if(position == word.size()){
        return;
    }
    TrieNode** child = &current_node->children[character_mapping[word[position]]];
    remove_word_recursive(*child, word, position + 1);
    (*child)->words_touching--;
    if(position == word.size() - 1){
        (*child)->words_ending--;
    }
    if((*child)->words_touching == 0){
        delete (*child);
        *child = NULL;
    }
}

int Trie::count_word(std::string &word){
    TrieNode *current_node = root;
    for(int i = 0;i < word.size();i++){
        TrieNode** child = &current_node->children[character_mapping[word[i]]];
        if((*child) == NULL){
            return 0;
        }
        if(i == word.size() - 1){
            return (*child)->words_ending;
        }
        current_node = *child;
    }
    return 0;
}

int Trie::longest_prefix(std::string &word){
    TrieNode *current_node = root;
    for(int i = 0;i < word.size();i++){
        TrieNode** child = &current_node->children[character_mapping[word[i]]];
        if((*child) == NULL){
            return i;
        }
        current_node = *child;
    }
    return word.size();
}

char s[22];

int main()
{
    std::unordered_map <char, short> mp;
    for(int i = 0;i < 26;i++){
        char ch = (char)i + 'a';
        mp[ch] = i;
    }
    Trie trie(26, mp);
    freopen("trie.in", "r", stdin);
    freopen("trie.out", "w", stdout);
    int op;
    while(scanf("%d", &op) == 1){
        scanf("%s", s);
        string s_new(s);
        if(op == 0){
            trie.add_word(s_new);
        }else if(op == 1){
            trie.remove_word(s_new);
        }else if(op == 2){
            printf("%d\n", trie.count_word(s_new));
        }else{
            printf("%d\n", trie.longest_prefix(s_new));
        }
    }
    return 0;
}