Cod sursa(job #1307181)

Utilizator gabrieligabrieli gabrieli Data 1 ianuarie 2015 15:23:02
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.71 kb
#include <fstream>
#include <string>
#include <memory>
#include <stack>
#include <iterator>
#include <iostream>
using namespace std;

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

class Trie {
public:
    Trie() : root {nullptr}
    {}
    
    void add(const string& word);
    int count(const string& word) const;
    void remove(const string& word);
    int longestPrefix(const string& word) const;
    
    ~Trie() {
        if (root) delete root;
    }
private:
    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 = new Node();
    
    auto node = root;
    for (auto ch = word.begin(); ch != word.end(); ++ch) {
        if (!node->next[*ch - 'a'])
            node->next[*ch - 'a'] = new 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< 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();
    
    for (auto ch = word.rbegin(); del && ch != word.rend(); ++ch) {
        node = st.top(); st.pop();
        delete node->next[*ch - 'a'];
        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();
}