Cod sursa(job #1346149)

Utilizator thewildnathNathan Wildenberg thewildnath Data 18 februarie 2015 02:30:12
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.31 kb
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <cstring>

using namespace std;

const int LET = 26;

struct node
{
  int count, sum, father;
  int trans[LET];

  node (int father = 0)
  {
    this->father = father;
    this->count = 0;
    this->sum = 0;
    memset (trans, 0, sizeof (trans));
  }
};

class trie
{
private:
  vector <node> states;

  int addState (int father = 0)
  {
    states.push_back (node (father));
    return states.size () - 1;
  }

  pair<int, int> getPos (const string &str)
  {
    int st = 0, length = 0;
    for (int i = 0; i < str.size (); ++i, ++length)
    {
      int let = str[i] - 'a';
      if (states[st].trans[let] && states[states[st].trans[let]].sum != 0)
        st = states[st].trans[let];
      else
        return make_pair(-1, length);
    }
    return make_pair(st, length);
  }
public:
  trie ()
  {
    addState ();
  }

  void addWord (const string &str)
  {
    int st = 0;
    for (int i = 0; i < str.size (); ++i)
    {
      ++states[st].sum;
      int let = str[i] - 'a';
      if (states[st].trans[let] == 0)
      {
        int aux = addState (st);
        states[st].trans[let] = aux;
      }
      st = states[st].trans[let];
    }
    ++states[st].count;
    ++states[st].sum;
  }

  void deleteWord (const string &str)
  {
    pair<int, int> pr = getPos (str);
    int st = pr.first;

    if (pr.second == str.size ())
    {
      --states[st].count;
      while (st != 0)
      {
        --states[st].sum;
        st = states[st].father;
      }
    }
  }

  int query (const string &str)
  {
    pair<int, int> pr = getPos (str);
    int st = pr.first;

    if (pr.second == str.size ())
      return states[st].count;
    else
      return 0;
  }

  int prefix (const string &str)
  {
    pair<int, int> pr = getPos (str);
    return pr.second;
  }
};

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;
}