Cod sursa(job #2249513)

Utilizator cosmin.pascaruPascaru Cosmin cosmin.pascaru Data 30 septembrie 2018 00:02:22
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 4.41 kb
#include <stdio.h>
#include <iostream>
#include <fstream>
#include <vector>
#include <array>
#include <algorithm>
#include <vector>
#include <stack>
#include <set>
#include <assert.h>
#include <queue>
#include <chrono>
#include <cstring>
#include <memory>
 
using LL = long long;
using ULL = int long long;
 
const std::string _problemName = "trie";
 
namespace std {
std::ifstream fin(_problemName + ".in");
std::ofstream fout(_problemName + ".out");
}
 
#define USE_FILES
 
#ifdef USE_FILES
#define cin fin
#define cout fout
#endif

struct TrieNode {

public:

    TrieNode() {
        memset(this, 0, sizeof(*this));
    }

    void insert(const std::string& word) {
        insertImpl(word, 0);
    }

    bool erase(const std::string& word) {
        return eraseImpl(word, 0);
    }

    int count(const std::string& word) {
        return countImpl(word, 0);
    }

    int longestPrefix(const std::string& word) {
        return longestPrefixImpl(word, 0);
    }

private:
    void insertImpl(const std::string& word, int pos) {
        if (pos == word.size()) {
            ++_count;
            return;
        }

        char nextChar = word[pos];
        getOrCreateNext(nextChar).insertImpl(word, pos + 1);
    }

    bool eraseImpl(const std::string& word, int pos) {
        if (pos == word.size()) {
            if (_count <= 0) {
                return false;
            }

            --_count;
            return true;
        }

        char nextChar = word[pos];

        if (!hasNext(nextChar)) {
            return false;
        }
    
        bool ok = getNext(nextChar).eraseImpl(word, pos + 1);

        if (ok && getNext(nextChar).isEmpty()) {
            deleteNext(nextChar);
        }

        return ok;
    }

    int countImpl(const std::string& word, int pos) {
        if (pos >= word.size()) {
            return _count;
        }

        char nextChar = word[pos];
        
        if (!hasNext(nextChar)) {
            return 0;
        }
        return getNext(nextChar).countImpl(word, pos + 1);
    }

    int longestPrefixImpl(const std::string& word, int pos) {
        if (pos >= word.size()) {
            return pos;
        }
        
        char nextChar = word[pos];
        
        if (!hasNext(nextChar)) {
            return pos;
        }
        return getNext(nextChar).longestPrefixImpl(word, pos + 1);
    }

    TrieNode& getOrCreateNext(char nextChar) {
        if (!hasNext(nextChar)) {
            getNextPtr(nextChar) = allocNode();
        }

        return getNext(nextChar);
    }

    bool isEmpty() {
        return (_count == 0 && _childrenCount == 0);
    }

    bool hasNext(char nextChar) {
        return (nullptr != getNextPtr(nextChar));
    }

    TrieNode& getNext(char nextChar) {
        return *getNextPtr(nextChar);
    }

    TrieNode*& getNextPtr(char nextChar) {
        return _next[nextChar - 'a'];
    }

    TrieNode* allocNode() {
        ++_childrenCount;
        return new TrieNode();
    }
    
    void deleteNode(TrieNode*& node) {
        --_childrenCount;
        delete node;
        node = nullptr;
    }

    void deleteNext(char c) {
        TrieNode*& nextNode = getNextPtr(c);
        deleteNode(nextNode);
    }

private:
    
    int _count;

    TrieNode* _next[26];
    int _childrenCount;

};

struct Trie {

public:

    Trie() : _root (std::make_unique<TrieNode>()) {}

    void insert(const std::string& word) {
        _root->insert(word);
    }

    bool erase(const std::string& word) {
        return _root->erase(word);
    }

    int count(const std::string& word) {
        return _root->count(word);
    }

    int longestPrefix(const std::string& word) {
        return _root->longestPrefix(word);
    }

private:

    std::unique_ptr<TrieNode> _root;
};

int main() {

    int op;
    std::string word;

    Trie trie;

    while (std::cin >> op >> word) {
        switch (op) {
            case 0:
                trie.insert(word);
                break;
            case 1:
                trie.erase(word);
                break;
            case 2:
                std::cout << trie.count(word) << '\n';
                break;
            case 3:
                std::cout << trie.longestPrefix(word) << '\n';
                break;
        }
    }

    return 0;
}