Cod sursa(job #1592907)

Utilizator dorumusuroiFMI - Doru Musuroi dorumusuroi Data 8 februarie 2016 09:25:54
Problema Trie Scor 55
Compilator cpp Status done
Runda Arhiva educationala Marime 2.94 kb
#include <bits/stdc++.h>
using namespace std;

class Trie{
private:
    int words;
    int prefixes;
    Trie* edges[26];
    void addWordWorker(Trie *vertex, string::iterator begin, string::iterator end){
        if(end - begin == 0){
            vertex->words += 1;
        }
        else{
            vertex->prefixes += 1;
            char c = *begin;
            if(vertex->edges[c-'a'] == NULL)
                vertex->edges[c-'a'] = new Trie();
            addWordWorker(vertex->edges[c-'a'],begin+1, end);
        }
    }
    void removeWordWorker(Trie *vertex, string::iterator begin, string::iterator end){
        if(end - begin == 0){
            vertex->words -= 1;
        }
        else{
            char c = *begin;
            vertex->prefixes -= 1;
            removeWordWorker(vertex->edges[c-'a'], begin+1, end);
        }
    }
    int countWordsWorker(Trie *vertex,string::iterator begin, string::iterator end){
        if(end - begin == 0){
            return vertex->words;
        }
        else{
            char c = *begin;
            if(vertex->edges[c-'a'] == NULL)
                return 0;
            else
                return countWordsWorker(vertex->edges[c-'a'], begin+1, end);
        }
    }
    string::iterator countPrefixWorker(Trie *vertex, string::iterator begin, string::iterator end){
        if(end - begin == 0)
            return begin;
        else{
            if(vertex->prefixes == 0 && vertex->words == 0)
                return begin-1;
            char c = *begin;
            if(vertex->edges[c-'a'] == NULL)
                return begin;
            else
                return countPrefixWorker(vertex->edges[c-'a'], begin+1, end);
        }
    }
public:
    Trie(){
        this->words = 0;
        this->prefixes = 0;
        for(int i = 0; i < 26; ++i)
            this->edges[i] = NULL;
    }
    void addWord(string word){
        addWordWorker(this, word.begin(), word.end());
    }
    void removeWord(string word){
        removeWordWorker(this, word.begin(), word.end());
    }
    int countWords(string word){
        return countWordsWorker(this, word.begin(), word.end());
    }
    int countPrefix(string word){
        return countPrefixWorker(this, word.begin(), word.end()) - word.begin();
    }
};

const char iname[] = "trie.in";
const char oname[] = "trie.out";
const int INSERT = 0, DELETE = 1, GET_WORDS = 2, GET_PREFIXES = 3;


int main()
{
    Trie trie;
    ifstream in(iname);
    ofstream out(oname);
    int op;
    string word;
    while(in >> op >> word){
        if(op == INSERT){
            trie.addWord(word);
        }
        else if(op == DELETE){
            trie.removeWord(word);
        }
        else if(op == GET_WORDS){
            out << trie.countWords(word) << '\n';
        }
        else if(op == GET_PREFIXES){
            out << trie.countPrefix(word) << '\n';
        }
    }
    return 0;
}