Cod sursa(job #1346157)

Utilizator thewildnathNathan Wildenberg thewildnath Data 18 februarie 2015 03:07:11
Problema Trie Scor 5
Compilator cpp Status done
Runda Arhiva educationala Marime 2.27 kb
#include <fstream>
#include <iostream>
#include <string>
#include <vector>

using namespace std;

const int LET = 26;

struct node
{
  int count, level, sons, letter;
  node* trans[LET];
  node* father;

  node (node* father = NULL, int letter = -1)
  {
    if (father == NULL)
        this->level = 0;
    else
        this->level = father->level + 1, ++father->sons;
    this->count = 0;
    this->father = father;
    this->sons = 0;
    this->letter = letter;
    for (int i =0; i < LET; ++i)
      trans[i] = NULL;
  }
};

class trie
{
private:
  node* root;

  node* addState (node* father = NULL, int letter = -1)
  {
    return new node (father, letter);
  }

  node* getPos (const string &str)
  {
    node* st = root;
    for (int i = 0; i < str.size (); ++i)
    {
      int let = str[i] - 'a';
      if (st->trans[let] != NULL)
        st = st->trans[let];
      else
        break;
    }
    return st;
  }
public:
  trie ()
  {
    root = addState ();
  }

  void addWord (const string &str)
  {
    node* st = root;
    for (int i = 0; i < str.size (); ++i)
    {
      int let = str[i] - 'a';
      if (st->trans[let] == NULL)
        st->trans[let] = addState (st, let);
      st = st->trans[let];
    }
    ++st->count;
  }

  void deleteWord (const string &str)
  {
    node* st = getPos (str);

    if (st->level == str.size ())
    {
      --st->count;
      while (st != root && st->sons == 0)
      {
        node* father = st->father;
        father->trans[st->letter] = NULL;
        delete[] st;
        st = father;
      }
    }
  }

  int query (const string &str)
  {
    node* st = getPos (str);

    if (st->level == str.size ())
      return st->count;
    else
      return 0;
  }

  int prefix (const string &str)
  {
    node* st = getPos (str);
    return st->level;
  }
};

trie tr;

int main ()
{
  ifstream cin("trie.in");
  ofstream cout("trie.out");
  int type, x;
  string str;

  while (cin >> type >> str)
  {
    if (type == 0)
      tr.addWord (str);
    else if (type == 1)
      tr.deleteWord (str);
    else if (type == 2)
      cout << tr.query (str) << '\n';
    else if (type == 3)
      cout << tr.prefix (str) << '\n';
  }

  return 0;
}