Cod sursa(job #2561678)

Utilizator memecoinMeme Coin memecoin Data 29 februarie 2020 08:05:02
Problema Trie Scor 55
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.43 kb
#include <fstream>
#include <iomanip>
#include <string>
#include <stdio.h>
#include <vector>
#include <algorithm>
#include <math.h>
#include <set>
#include <map>
#include <string.h>
#include <queue>
#include <stack>

#define INF 0x3f3f3f3f

using namespace std;

#ifdef DEBUG
string name = "data";
#else
string name = "trie";
#endif

ifstream fin(name + ".in");
ofstream fout(name + ".out");

class Trie {
public:
    Trie *children[26];
    
    int nrEnding = 0;
    int subTreeWordCount = 0;
    
    Trie() {
        memset(children, 0, sizeof(children));
    }
    
    void insert(string &s, int k) {
        if (k >= s.size()) {
            nrEnding++;
            return;
        }
        
        char c = s[k] - 'a';
        
        if (children[c] == NULL) {
            children[c] = new Trie();
        }
        
        subTreeWordCount++;
        children[c]->insert(s, k + 1);
    }
    
    void remove(string &s, int k, bool &didRemove) {
        if (k >= s.size()) {
            nrEnding--;
            return;
        }
        
        char c = s[k] - 'a';
        
        if (children[c] == NULL) {
            return;
        }
        
        children[c]->remove(s, k + 1, didRemove);
        if (didRemove) {
            subTreeWordCount--;
        }
    }
    
    int count(string &s, int k) {
        if (k == s.size()) {
            return nrEnding;
        }
        
        char c = s[k] - 'a';
        
        if (children[c] == NULL) {
            return 0;
        }
    
        return children[c]->count(s, k + 1);
    }
    
    int longestPrefix(string &s, int k) {
        
        char c = s[k] - 'a';
        
        if (k == s.size() || children[c] == NULL) {
            return k;
        }
        
        if (children[c]->nrEnding == 0 && children[c]->subTreeWordCount == 0) {
            return k;
        }
        
        return children[c]->longestPrefix(s, k + 1);
    }
};

Trie trie;

int main() {
    
    while (true) {
        int op;
        string s;
        
        fin >> op >> s;
        
        if (fin.eof()) {
            break;
        }
        
        if (op == 0) {
            trie.insert(s, 0);
        }
        if (op == 1) {
            bool didRemove = false;
            trie.remove(s, 0, didRemove);
        }
        if (op == 2) {
            fout << trie.count(s, 0) << "\n";
        }
        if (op == 3) {
            fout << trie.longestPrefix(s, 0) << "\n";
        }
    }
    
    return 0;
}