Cod sursa(job #2462746)

Utilizator preda.andreiPreda Andrei preda.andrei Data 27 septembrie 2019 19:36:58
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.35 kb
#include <fstream>
#include <string>
#include <unordered_map>

using namespace std;

class Node
{
public:
    bool HasSon(char ch) const
    {
        auto it = sons_.find(ch);
        return it != sons_.end() && it->second->count_ > 0;
    }

    Node* GetSon(char ch)
    {
        auto it = sons_.find(ch);
        if (it != sons_.end()) {
            return it->second;
        }
        return sons_[ch] = new Node();
    }

    void Visit(bool end = false)
    {
        count_ += 1;
        if (end) {
            end_count_ += 1;
        }
    }

    void Leave(bool end = false)
    {
        count_ = max(0, count_ - 1);
        if (end) {
            end_count_ = max(0, end_count_ - 1);
        }
    }

    int Count() const { return count_; }
    int EndCount() const { return end_count_; }

private:
    unordered_map<char, Node*> sons_;
    int count_ = 0;
    int end_count_ = 0;
};

class Trie
{
public:
    void Add(const string &str);
    void Remove(const string &str);

    int Count(const string &str) const;
    int MaxPrefix(const string &str) const;

private:
    Node *root_ = new Node();
};

void Trie::Add(const string &str)
{
    auto node = root_;
    for (const auto &ch : str) {
        node->Visit();
        node = node->GetSon(ch);
    }
    node->Visit(true);
}

void Trie::Remove(const string &str)
{
    auto node = root_;
    for (const auto &ch : str) {
        node->Leave();
        node = node->GetSon(ch);
    }
    node->Leave(true);
}

int Trie::Count(const string &str) const
{
    auto node = root_;
    for (const auto &ch : str) {
        if (!node->HasSon(ch)) {
            return 0;
        }
        node = node->GetSon(ch);
    }
    return node->EndCount();
}

int Trie::MaxPrefix(const string &str) const
{
    auto node = root_;
    auto len = 0;
    for (const auto &ch : str) {
        if (!node->HasSon(ch)) {
            return len;
        }
        len += 1;
        node = node->GetSon(ch);
    }
    return len;
}

int main()
{
    ifstream fin("trie.in");
    ofstream fout("trie.out");

    int type;
    string str;

    Trie trie;
    while (fin >> type >> str) {
        switch (type) {
            case 0: trie.Add(str); break;
            case 1: trie.Remove(str); break;
            case 2: fout << trie.Count(str) << "\n"; break;
            case 3: fout << trie.MaxPrefix(str) << "\n"; break;
        }
    }
    return 0;
}