Cod sursa(job #2814772)

Utilizator Asgari_ArminArmin Asgari Asgari_Armin Data 8 decembrie 2021 16:30:15
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.92 kb
#include <bits/stdc++.h>

using namespace std;

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

const int SIGMA = 26;

struct Node {
  int full_words = 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->full_words;
  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->full_words;
  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->full_words;
  return query_(1 + s, node->v[s[0] - 'a']);
}

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