Cod sursa(job #2916901)

Utilizator CiobaCatalinCioba Catalin CiobaCatalin Data 2 august 2022 11:45:46
Problema Trie Scor 70
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.85 kb
#include <fstream>
#include <iostream>
#include <unordered_map>

using namespace std;

class TrieNode
{
public:
    TrieNode();
    void add(string word);
    void remove(string word);
    int count(string word);
    int prefix(string word);
    void print();
    ~TrieNode();

private:
    int endingHere;
    int prefixCount;
    unordered_map<char, TrieNode *> children;
};

int main()
{
    ios_base::sync_with_stdio(false);
    freopen("trie.in", "r", stdin);
    freopen("trie.out", "w", stdout);

    int query;
    string word;

    TrieNode trie;

    while (cin >> query)
    {
        cin.ignore(1);
        getline(cin, word);

        if (query == 0)
        {
            trie.add(word);
        }
        else if (query == 1)
        {
            trie.remove(word);
        }
        else if (query == 2)
        {
            cout << trie.count(word) << '\n';
        }
        else
        {
            cout << trie.prefix(word) << '\n';
        }
    }

    return 0;
}

TrieNode::TrieNode()
{
    this->endingHere = 0;
    this->prefixCount = 0;
    this->children = {};
}

void TrieNode::add(string word)
{
    auto current = this;

    for (auto it = word.begin(); it != word.end(); ++it)
    {
        // cout << "Looking for " << *it << endl;
        auto nodeCount = current->children.count(*it);
        if (nodeCount == 0)
        {
            current->children[*it] = new TrieNode();
        }
        auto node = current->children[*it];
        current->prefixCount++;
        current = node;
    }

    current->prefixCount++;
    current->endingHere++;

    // cout << word << " -> " << current->endingHere << endl;
};

void TrieNode::remove(string word)
{
    auto current = this;

    for (auto it = word.begin(); it != word.end(); ++it)
    {
        current->prefixCount--;
        current = current->children[*it];
    }

    current->prefixCount--;
    current->endingHere--;

    // cout << word << " -> " << current->endingHere << endl;
}

int TrieNode::count(string word)
{
    auto current = this;

    for (auto it = word.begin(); it != word.end(); ++it)
    {
        auto node = current->children.find(*it);
        if (node != current->children.end() && node->second->prefixCount > 0)
        {
            current = node->second;
        }
        else
        {
            return 0;
        }
    }

    return current->endingHere;
}

int TrieNode::prefix(string word)
{
    auto current = this;
    int prefix = 0;

    for (auto it = word.begin(); it != word.end(); ++it)
    {
        auto node = current->children.find(*it);
        if (node != current->children.end() && node->second->prefixCount > 0)
        {
            prefix++;
            current = node->second;
        }
        else
        {
            break;
        }
    }

    return prefix;
}

TrieNode::~TrieNode()
{
    for (auto it : this->children)
    {
        delete it.second;
    }
}