Cod sursa(job #2884043)

Utilizator PushkinPetolea Cosmin Pushkin Data 2 aprilie 2022 12:17:19
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 3.33 kb
#include <bits/stdc++.h>
using namespace std;

class Trie {
    struct TrieNode {
        int val;
        int next_count;
        vector<TrieNode*> next;

        TrieNode() {
            val = 0;
            next_count = 0;
            next.resize(26, NULL);
        }
    };

    TrieNode *root;

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

    void add_word(string word) {
        TrieNode *node = root;
        size_t i = 0;

        while(i < word.size()) {
            if(!node->next[word[i] - 'a']) {
                node->next[word[i] - 'a'] = new TrieNode();
                node->next_count++;
            }

            node = node->next[word[i] - 'a'];
            i++;
        }

        node->val++;
    }

    void dealloc_word(string word, TrieNode* node, size_t i) {
        if(i == word.size()) {
            return;
        }

        dealloc_word(word, node->next[word[i] - 'a'], i + 1);

        if(node->next[word[i] - 'a'] &&
           !node->next[word[i] - 'a']->val &&
           !node->next[word[i] - 'a']->next_count) {
            delete node->next[word[i] - 'a'];
            node->next[word[i] - 'a'] = NULL;
            node->next_count--;
        }
    }
    void erase_word(string word) {
        TrieNode *node = root;
        size_t i = 0;

        while(node && i < word.size()) {
            node = node->next[word[i] - 'a'];
            i++;
        }

        if(node)
            node->val--;

        dealloc_word(word, root, 0);
    }

    int get_count(string word) {
        TrieNode *node = root;
        size_t i = 0;

        while(node && i < word.size()) {
            node = node->next[word[i] - 'a'];
            i++;
        }

        if(node)
            return node->val;

        return 0;
    }

    int get_longest_prefix(string word) {
        TrieNode *node = root;
        size_t i = 0;

        while(node && i < word.size()) {
            node = node->next[word[i] - 'a'];
            i++;
        }

        if(!node)
            i--;

        return i;
    }

    void print() {
        queue<pair<TrieNode*, string>> q;

        q.push(make_pair(root, ""));
        while(q.size()) {
            auto node = q.front();
            q.pop();

            cout << "\"" << node.second << "\" -> ";
            cout << "val(" << node.first->val << ") ";
            cout << "next_count(" << node.first->next_count << ")\n";

            for(size_t i = 0; i < node.first->next.size(); i++)
                if(node.first->next[i])
                    q.push(make_pair(node.first->next[i], node.second + char(i + 'a')));
        }
    }
};

void solve(string filename_in, string filename_out) {
    ifstream in(filename_in);
    ofstream out(filename_out);

    Trie trie;

    int op;
    string word;
    while(in >> op >> word) {
        switch(op) {
            case 0: trie.add_word(word); break;
            case 1: trie.erase_word(word); break;
            case 2: out << trie.get_count(word) << '\n'; break;
            case 3: out << trie.get_longest_prefix(word) << '\n'; break;
            default: cout << "Unknown op: " << op << '\n';
        }

        //trie.print();
    }

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

int main() {
    solve("trie.in", "trie.out");

    return 0;
}