Cod sursa(job #1307173)

Utilizator gabrieligabrieli gabrieli Data 1 ianuarie 2015 14:57:55
Problema Trie Scor 80
Compilator cpp Status done
Runda Arhiva educationala Marime 2.46 kb
#include <fstream>
#include <string>
#include <memory>
#include <stack>
#include <iterator>
#include <iostream>
using namespace std;

struct Node {
    int nr;
    shared_ptr<Node> next[26];
    
    bool isLeaf()
    {
        for (int i = 0; i < 26; ++i) if (next[i])
            return false;
        return true;
    }
};

class Trie {
public:
    Trie()
    {}
    
    void add(const string& word);
    int count(const string& word) const;
    void remove(const string& word);
    int longestPrefix(const string& word) const;

private:
    shared_ptr<Node> root;
};

int main() {
    ifstream fin("trie.in");
    ofstream fout("trie.out");
    
    int op;
    string word;
    Trie trie;
    
    while (fin >> op >> word) {
        if (op == 0)
            trie.add(word);
        else if (op == 1)
            trie.remove(word);
        else if (op == 2)
            fout << trie.count(word) << '\n';
        else
            fout << trie.longestPrefix(word) << '\n';
    }
    
    return 0;
}

void Trie::add(const string& word) {
    if (!root) root = make_shared<Node>();
    
    auto node = root;
    for (auto ch = word.begin(); ch != word.end(); ++ch) {
        if (!node->next[*ch - 'a'])
            node->next[*ch - 'a'] = make_shared<Node>();
        node = node->next[*ch - 'a'];
    }
    node->nr += 1;
}

int Trie::count(const string& word) const {
    if (!root) return 0;
    
    auto node = root;
    for (auto ch = word.begin(); ch != word.end(); ++ch) {
        if (!node->next[*ch - 'a']) return 0;
        node = node->next[*ch - 'a'];
    }
    
    return node->nr;
}

void Trie::remove(const string& word) {
    if (!root) return;
    
    auto node = root;
    stack< shared_ptr<Node> > st;
    for (auto ch = word.begin(); ch != word.end(); ++ch) {
        if (!node->next[*ch - 'a']) return;  
        st.push(node);
        node = node->next[*ch - 'a'];
    }
    
    node->nr -= 1;
    bool del = (node->nr == 0 && node->isLeaf()) ? true : false;
    
    for (auto ch = word.rbegin(); del && ch != word.rend(); ++ch) {
        node = st.top(); st.pop();
        node->next[*ch - 'a'] = nullptr;
        
        del = (node->nr == 0) && node->isLeaf();
    }
}

int Trie::longestPrefix(const string& word) const {
    if (!root) return 0;
    
    auto node = root;
    for (auto ch = word.begin(); ch != word.end(); ++ch) {
        if (!node->next[*ch - 'a']) return distance(word.begin(), ch);
        node = node->next[*ch - 'a'];
    }
    
    return word.size();
}