Cod sursa(job #2760650)

Utilizator cristi_macoveiMacovei Cristian cristi_macovei Data 28 iunie 2021 15:43:17
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.16 kb
#include <iostream>
#include <fstream>
#include <queue>

const int SIGMA = 26;

class TrieNode {
public:
  TrieNode *parent;
  int value;
  int cnt_words;
  char letter;
  TrieNode* children[SIGMA];

  TrieNode(char letter, TrieNode* parent = nullptr, int value = 0) {
    this->letter = letter;
    this->value = value;
    this->parent = parent;
    this->cnt_words = 0;

    for (int i = 0; i < SIGMA; ++i)
      this->children[i] = nullptr;
  }
};

class Trie {
public:
  TrieNode *root;

  Trie() {
    this->root = new TrieNode('-');
  }

  ~Trie() {
    destroy(root);
  }

  void destroy(TrieNode* node) {
    for (int i = 0; i < SIGMA; ++i) {
      if (node->children[i] != nullptr) {
        destroy(node->children[i]);
        node->children[i] = nullptr;
      }
    }

    if (node != root)
      node->parent->children[node->letter - 'a'] = nullptr;

    delete node;
  }

  void dfs_debug(TrieNode* node, int depth) {
    if (depth)
      std::printf("node letter = %c, value = %d, words = %d, depth = %d, parent letter = %c\n", node->letter, node->value, node->cnt_words, depth, node->parent->letter);

    for (int i = 0; i < SIGMA; ++i) {
      if (node->children[i] != nullptr)
        dfs_debug(node->children[i], depth + 1);
    }
  }

  void add_word(const std::string& word) {
    TrieNode *node = root;
    for (char c: word) {
      if (node->children[c - 'a'] == nullptr)
        node->children[c - 'a'] = new TrieNode(c, node);

      ++node->children[c - 'a']->cnt_words;

      node = node->children[c - 'a'];
    }

    ++node->value;
  }

  int count(const std::string& word) {
    TrieNode *node = root;
    for (char c: word) {
      if (node->children[c - 'a'] == nullptr)
        return 0;

      node = node->children[c - 'a'];
    }

    return node->value;
  }

  void del(const std::string& word) {
    TrieNode *node = root;

    for (char c: word)
      node = node->children[c - 'a'];

    --node->value;

    TrieNode *del_from = nullptr;
    while (node != root) {
      --node->cnt_words;

      if (node->cnt_words == 0)
        del_from = node;

      node = node->parent;
    }

    if (del_from != nullptr)
      destroy(del_from);
  }

  int prefix(const std::string& word) {
//    std::cout << "prefix for " << word << '\n';
    TrieNode *node = root;

    int ans = 0;
    for (char c: word) {
      node = node->children[c - 'a'];

      if (node == nullptr)
        break;

//      std::cout << "++ for char " << c << '\n';
      ++ans;
    }

//    std::cout << "returns " << ans << '\n';
    return ans;
  }
};

int main() {
  std::ifstream in("trie.in");
  std::ofstream out("trie.out");

  Trie trie;

  int type;
  std::string word;
  while (in >> type >> word) {
    switch (type) {
      case 0:
        trie.add_word(word);
        break;

      case 1:
        trie.del(word);
        break;

      case 2:
        out << trie.count(word) << '\n';
        break;

      case 3:
        out << trie.prefix(word) << '\n';
        break;

      default:
        throw std::runtime_error("Unexpected operation type '" + std::to_string(type) + "'");
    }

//    trie.dfs_debug(trie.root, 0);
//    std::cout << '\n' << '\n';
  }

//  trie.dfs_debug(trie.root, 0);

  return 0;
}