Cod sursa(job #2814736)

Utilizator livliviLivia Magureanu livlivi Data 8 decembrie 2021 15:38:44
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.1 kb
//#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

ifstream cin("trie.in");
ofstream cout("trie.out");

struct Node {
    int occ = 0;
    int full_words = 0;
    vector<Node*> sons;

    Node() {
        sons.resize(26, nullptr);
    }
};

class Trie {
public:
    Trie() { root_ = nullptr; }

    void insert(char* s) { root_ = insert_(s, root_); }
    void erase(char* s) { root_ = erase_(s, root_); }
    int query(char* s) { return query_(s, root_); }
    int prefix(char* s) { return prefix_(s, root_); }

private:
    Node* root_;

    Node* insert_(char* s, Node* node);
    Node* erase_(char* s, Node* node);
    int query_(char* s, Node* node);
    int prefix_(char* s, Node* node);
};

Node* Trie::insert_(char *s, Node *node) {
    if (node == nullptr) {
        node = new Node();
    }

    node->occ++;
    if (s[0] == '\0') {
        node->full_words++;
    } else {
        node->sons[s[0] - 'a'] = insert_(s + 1, node->sons[s[0] - 'a']);
    }

    return node;
}

Node* Trie::erase_(char *s, Node *node) {
    if (node == nullptr) { return node; }

    node->occ--;
    if (s[0] == '\0') {
        node->full_words--;
    } else {
        node->sons[s[0] - 'a'] = erase_(s + 1, node->sons[s[0] - 'a']);
    }

    if (node->occ == 0) {
        delete node;
        node = nullptr;
    }

    return node;
}

int Trie::query_(char *s, Node *node) {
    if (node == nullptr) { return 0; }
    if (s[0] == '\0') { return node->full_words; }
    return query_(s + 1, node->sons[s[0] - 'a']);
}

int Trie::prefix_(char *s, Node *node) {
    if (node == nullptr || s[0] == '\0') { return 0; }
    if (node->sons[s[0] - 'a'] == nullptr) {
        return 0;
    } else {
        return 1 + prefix_(s + 1, node->sons[s[0] - 'a']);
    }
}

int main() {
    Trie trie;
    int op;
    char s[21];

    while (cin >> op) {
        cin >> s;
        if (op == 0) {
            trie.insert(s);
        } else if (op == 1) {
            trie.erase(s);
        } else if (op == 2) {
            cout << trie.query(s) << '\n';
        } else {
            cout << trie.prefix(s) << '\n';
        }
    }

    return 0;
}