Cod sursa(job #1110273)

Utilizator a_h1926Heidelbacher Andrei a_h1926 Data 17 februarie 2014 21:57:38
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.23 kb
#include <fstream>
#include <algorithm>
#include <map>

using namespace std;

class Trie {
  public:
    Trie():
      root(new Node()) {}

    void Insert(const string &word) {
        Node *node = root;
        for (int i = 0; i < int(word.length()); ++i) {
            ++node->size;
            if (node->sons.count(word[i]) == 0)
                node->sons[word[i]] = new Node();
            node = node->sons[word[i]];
        }
        ++node->size;
        ++node->count;
    }

    void Erase(const string &word) {
        Erase(root, word, 0);
    }

    int Count(const string &word) const {
        Node *node = root;
        for (int i = 0; i < int(word.length()); ++i) {
            if (node->sons.count(word[i]) == 0)
                return 0;
            node = node->sons[word[i]];
        }
        return node->count;
    }

    string LCP(const string &word) const {
        int i = 0;
        for (Node *node = root; i < int(word.length()) && node->sons.count(word[i]) > 0; node = node->sons[word[i++]]);
        return word.substr(0, i);
    }

  private:
    class Node {
      public:
        int size, count;
        map<char, Node*> sons;

        Node():
          size(0),
          count(0),
          sons(map<char, Node*>()) {}
    };

    Node *root;

    void Erase(Node *node, const string &word, const int index) {
        --node->size;
        if (index == int(word.length())) {
            --node->count;
            return;
        }
        Erase(node->sons[word[index]], word, index + 1);
        if (node->sons[word[index]]->size == 0) {
            delete node->sons[word[index]];
            node->sons.erase(node->sons.find(word[index]));
        }
    }
};

int main() {
    ifstream cin("trie.in");
    ofstream cout("trie.out");
    int type;
    string word;
    Trie tree = Trie();
    while (cin >> type >> word) {
        if (type == 0)
            tree.Insert(word);
        else if (type == 1)
            tree.Erase(word);
        else if (type == 2)
            cout << tree.Count(word) << "\n";
        else if (type == 3)
            cout << int(tree.LCP(word).length()) << "\n";
    }
    cin.close();
    cout.close();
    return 0;
}