Cod sursa(job #2150817)

Utilizator preda.andreiPreda Andrei preda.andrei Data 3 martie 2018 19:55:43
Problema Trie Scor 50
Compilator cpp Status done
Runda Arhiva educationala Marime 2.36 kb
#include <fstream>
#include <map>
#include <string>

using namespace std;

class Trie
{
  struct Node
  {
    int count = 0;
    map<char, struct Node*> sons;
  };

public:
  Trie() : root(new Node) {}

  void Add(const string &s) { Add(root, s); }

  void Remove(const string &s, int count) { Remove(root, s, count); }

  int GetCount(const string &s) const { return GetCount(root, s); }

  int LongestPrefix(const string &s) const { return LongestPrefix(root, s); }

private:
  void Add(Node *node, const string &s, size_t pos = 0);

  bool Remove(Node *node, const string &s, int count, size_t pos = 0);

  int GetCount(Node *node, const string &s, size_t pos = 0) const;

  int LongestPrefix(Node *node, const string &s, size_t pos = 0) const;

  Node *root;
};

void Trie::Add(Node *node, const string &s, size_t pos)
{
  if (pos >= s.size()) {
    node->count += 1;
    return;
  }

  if (node->sons.count(s[pos]) == 0) {
    node->sons.insert({s[pos], new Node});
  }
  Add(node->sons[s[pos]], s, pos + 1);
}

bool Trie::Remove(Node *node, const string &s, int count, size_t pos)
{
  if (pos >= s.size()) {
    node->count -= count;
    if (node->count <= 0) {
      delete node;
      return true;
    }
    return false;
  }

  auto it = node->sons.find(s[pos]);
  if (it == node->sons.end()) {
    return false;
  }

  if (Remove(it->second, s, count, pos + 1)) {
    node->sons.erase(s[pos]);
  }
  if (node != root && node->sons.empty()) {
    delete node;
    return true;
  }
  return false;
}

int Trie::GetCount(Node *node, const string &s, size_t pos) const
{
  if (pos >= s.size()) {
    return node->count;
  }

  auto it = node->sons.find(s[pos]);
  if (it == node->sons.end()) {
    return 0;
  }
  return GetCount(it->second, s, pos + 1);
}

int Trie::LongestPrefix(Node *node, const string &s, size_t pos) const
{
  if (pos >= s.size()) {
    return pos;
  }

  auto it = node->sons.find(s[pos]);
  if (it == node->sons.end()) {
    return pos;
  }
  return LongestPrefix(it->second, s, pos + 1);
}

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

  Trie trie;
  int type;
  string word;

  while (fin >> type >> word) {
    if (type == 0) {
      trie.Add(word);
    } else if (type == 1) {
      trie.Remove(word, 1);
    } else if (type == 2) {
      fout << trie.GetCount(word) << "\n";
    } else {
      fout << trie.LongestPrefix(word) << "\n";
    }
  }
  return 0;
}