Cod sursa(job #3042008)

Utilizator Asgari_ArminArmin Asgari Asgari_Armin Data 3 aprilie 2023 15:14:52
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.97 kb
#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

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

const int SIGMA = 26;

struct Node {
    int fullWords = 0, freq = 0;
    vector <Node*> v;

    Node() : v(SIGMA, nullptr) {}
};

struct Trie {
    private: 
        Node* root_ = nullptr;
        Node* insert_(char *s, Node *node);
        Node* erase_(char *s, Node *node);
        int query_(char *s, Node *node);
        int prefix_(char *s, Node *node);

    public: 
        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_);};
};

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

    ++node->freq;
    if (s[0] == '\0')
        ++node->fullWords;
    else
        node->v[s[0] - 'a'] = insert_(1 + s, node->v[s[0] - 'a']);
    return node;
}       

Node *Trie::erase_(char *s, Node *node) {
    if (node == nullptr)
        return node;
    
    --node->freq;
    if (s[0] == '\0')
        --node->fullWords;
    else 
        node->v[s[0] - 'a'] = erase_(1 + s, node->v[s[0] - 'a']);
    
    if (node->freq == 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->fullWords;
    return query_(1 + s, node->v[s[0] - 'a']);
}

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

int main() {
  Trie trie;
  int op;
  char v[30];

  while( fin >> op >> v ){
    if( op == 0 )
      trie.insert(v);
    else if( op == 1 )
      trie.erase(v);
    else if( op == 2 )
      fout << trie.query(v) << "\n";
    else
      fout << trie.prefix(v) << "\n";
  }
  return 0;
}