Cod sursa(job #2916943)

Utilizator CiobaCatalinCioba Catalin CiobaCatalin Data 2 august 2022 14:53:46
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.82 kb
#include <fstream>
#include <iostream>
#include <map>
#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);

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

int main()
{
    // ios_base::sync_with_stdio(false);
    // cin.tie(0);
    // cout.tie(0);
    // freopen("trie.in", "r", stdin);
    // freopen("trie.out", "w", stdout);
    ifstream in("trie.in");
    ofstream out("trie.out");

    int query;
    string word;

    TrieNode trie;

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

        if (query == 0)
        {
            trie.add(word);
        }
        else if (query == 1)
        {
            trie.remove(word);
        }
        else if (query == 2)
        {
            out << trie.count(word) << '\n';
        }
        else
        {
            out << 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;
}