Cod sursa(job #2761267)

Utilizator cristi_macoveiMacovei Cristian cristi_macovei Data 1 iulie 2021 13:40:02
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.6 kb
#include <iostream>
#include <fstream>
#include <queue>

const int SIGMA = 26;

class Trie {
private:
  class TrieNode {
  public:
    int count;
    int childrenCount;
    TrieNode* parent;
    TrieNode* children[SIGMA];

    TrieNode(TrieNode* parent = nullptr) {
      this->count         = 0;
      this->childrenCount = 0;
      this->parent        = parent;

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

public:
  TrieNode* root;

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

  ~Trie() {
    destroy(root);
  }

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

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

    delete node;
  }

  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(node);
        ++node->childrenCount;
      }

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

    ++node->count;
  }

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

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

    --node->count;

    int letter_pos = word.size() - 1;
    while(node->childrenCount == 0 && node->count == 0 && node != root) {
      TrieNode* parent = node->parent;

      delete node;
      parent->children[word[letter_pos] - 'a'] = nullptr;
      --parent->childrenCount;

      node = parent;
      --letter_pos;
    }
  }

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

  int prefix(const std::string& word) {
    TrieNode* node = root;

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

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

    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) + "'");
    }
  }

  return 0;
}