Cod sursa(job #2300851)

Utilizator TooHappyMarchitan Teodor TooHappy Data 12 decembrie 2018 10:42:36
Problema Trie Scor 5
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.89 kb
#include <bits/stdc++.h>

using namespace std;

ifstream in("trie.in");
ofstream out("trie.out");

struct TrieNode {
    map< char, TrieNode* > children;
    int appearances;
    bool endOfWord;

    TrieNode () {
        appearances = 0;
        endOfWord = false;
    }
};

class Trie {
private:
    TrieNode *root;

    bool deleteWord (TrieNode *current, string word, int idx) {
        if (idx == (int)word.size()) {
            if (!current->endOfWord) {
                return false;
            }
            current->endOfWord = false;
            return ((int)current->children.size() == 0);
        }

        TrieNode *nextNode = current->children[word[idx]];
        if (nextNode == NULL) {
            return false;
        }

        bool shouldDeleteCurrentNode = deleteWord(nextNode, word, idx + 1);

        if (shouldDeleteCurrentNode) {
            current->children.erase(current->children.find(word[idx]));

            return ((int)current->children.size() == 0);
        }

        return false;
    }
public:
    Trie () {
        root = new TrieNode;
    }

    void insert (string word) {
        TrieNode *current = this->root;

        for (auto ch: word) {
            TrieNode *nextNode = current->children[ch];
            if (nextNode == NULL) {
                nextNode = new TrieNode;
                current->children[ch] = nextNode;
            }

            current = nextNode;
        }
        
        current->appearances++;
        current->endOfWord = true;
    }

    void deleteWord (string word) {
        deleteWord(this->root, word, 0);
    }

    int getAppearances (string word) {
        TrieNode *current = this->root;

        for (auto ch: word) {
            TrieNode *nextNode = current->children[ch];

            if (nextNode == NULL) {
                if(current->endOfWord) {
                    return current->appearances;
                } else {
                    return 0;
                }
            }

            current = nextNode;
        }

        return current->appearances;
    }

    int getLongestPrefix (string word) {
        TrieNode *current = this->root;

        int ans = 0;
        for (auto ch: word) {
            TrieNode *nextNode = current->children[ch];

            if (nextNode == NULL) {
                return ans;
            }

            ++ans;
            current = nextNode;
        }

        return ans;
    }
};

int main() {
    ios::sync_with_stdio(false); in.tie(0); out.tie(0);

    Trie *myTrie = new Trie;

    int op;
    string word;
    while (in >> op >> word) {
        if (op == 0) {
            myTrie->insert(word);
        }
        if (op == 1) {
            myTrie->deleteWord(word);
        }
        if (op == 2) {
            out << myTrie->getAppearances(word) << "\n";
        }
        if (op == 3) {
            out << myTrie->getLongestPrefix(word) << "\n";
        }
    }

    in.close(); out.close();

    return 0;
}