Cod sursa(job #2966532)

Utilizator DKMKDMatei Filibiu DKMKD Data 17 ianuarie 2023 19:36:55
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.99 kb
#include <bits/stdc++.h>
#include <unordered_map>
#include <map>

using namespace std;

ifstream fin("trie.in");
ofstream fout("trie.out");
//-----------------------------------------
#define ll long long
#define pii pair<int, int>
#define pll pair<long long, long long>
#define vi vector<int>
#define vll vector<long long>
#define mii map<int, int>
#define si set<int>
#define sc set<char>
//-----------------------------------------
#define f(i,s,e) for(int i=s;i<=e;++i)
#define nf(i,s,e) for(i=s;i<e;++i)
#define rf(i,e,s) for(i=e;i>=s;--i)
#define pb push_back
#define eb emplace_back
#define fi first
#define se second
#define sp <<" "
//------------------------------------------
const int NMAX = 2e1 + 5;
const int KMAX = 1e1 + 5;
const int MOD = 1e9 + 7;
const int INF = 0x3f3f3f3f;
//------------------------------------------
int query;
string cuv;
class Trie{
public:
    Trie();
    void add(string word);
    void remove(string word);
    int count(string word);
    int prefix(string word);

private:
    int end;
    int pref;
    map<char, Trie*> ch;
};
//------------------------------------------
void read() {
    Trie trie;
    while (fin >> query) {
        fin.ignore(1);
        getline(fin, cuv);
        if (query == 0)
            trie.add(cuv);
        else if (query == 1)
            trie.remove(cuv);
        else if (query == 2)
            fout << trie.count(cuv) << "\n";
        else
            fout << trie.prefix(cuv) << "\n";
    }
}
Trie::Trie(){
    this->end = 0;
    this->pref = 0;
    this->ch = {};
}
void Trie::add(string word){
    auto current = this;
    for (auto it = word.begin(); it != word.end(); ++it){
        auto nodeCount = current->ch.count(*it);
        if (nodeCount == 0)
            current->ch[*it] = new Trie();
        auto node = current->ch[*it];
        current->pref++;
        current = node;
    }
    current->pref++;
    current->end++;
};
void Trie::remove(string word){
    auto current = this;
    for (auto it = word.begin(); it != word.end(); ++it){
        current->pref--;
        current = current->ch[*it];
    }
    current->pref--;
    current->end--;
}

int Trie::count(string word){
    auto current = this;
    for (auto it = word.begin(); it != word.end(); ++it){
        auto node = current->ch.find(*it);
        if (node != current->ch.end() && node->second->pref > 0)
            current = node->second;
        else
            return 0;
    }
    return current->end;
}

int Trie::prefix(string word){
    auto current = this;
    int prefix = 0;
    for (auto it = word.begin(); it != word.end(); ++it){
        auto nod = current->ch.find(*it);
        if (nod != current->ch.end() && nod->second->pref > 0){
            prefix++;
            current = nod->second;
        }
        else
            break;
    }
    return prefix;
}
int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    read();
    return 0;
}