Cod sursa(job #2566273)

Utilizator SqueekDanielTodasca Daniel SqueekDaniel Data 2 martie 2020 20:10:56
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.89 kb
#include <bits/stdc++.h>

#define MAXN        200005
#define FILENAME    std::string("trie")

struct Node {
    Node() { for (int i=0; i<30; ++i) sons[i] = nullptr; freq = weight = 0; }
    Node *get(char ch) { return sons[ch-'a']; }
    Node *add(char ch) { return sons[ch-'a'] = new Node(); }
    void erase(char ch) { Node *node = sons[ch-'a']; sons[ch-'a'] = 0; delete node; }
    int freq, weight;

    Node *sons[30];
}   *trie;

void add(char *ptr, int len, Node *&node = trie) {
    ++ node->freq;
    if (len == 0) { ++ node->weight; return; }
    char ch = ptr[0];

    Node *down = node->get(ch);
    if (down == nullptr) down = node->add(ch);
    add(ptr+1, len-1, down);
}
void erase(char *ptr, int len, Node *&node = trie) {
    -- node->freq;
    if (len == 0) { -- node->weight; return; }
    char ch = ptr[0];

    Node *down = node->get(ch);
    if (down == nullptr) return;
    erase(ptr+1, len-1, down);
    if (down->freq == 0) node->erase(ch);
}
int search(char *ptr, int len, Node *&node = trie) {
    if (len == 0) return node->weight;
    char ch = ptr[0];

    Node *down = node->get(ch);
    if (down == nullptr) return 0;
    return search(ptr+1, len-1, down);
}
int prefix(char *ptr, int len, int depth = 0, Node *&node = trie) {
    if (len == 0) return depth;
    char ch = ptr[0];

    Node *down = node->get(ch);
    if (down == nullptr) return depth;
    return prefix(ptr+1, len-1, depth+1, down);
}

std::ifstream input (FILENAME+".in");
std::ofstream output(FILENAME+".out");

int main()
{
    int t;
    char str[MAXN];
    trie = new Node();
    while (input >> t >> str) {
        int len = strlen(str);
        if (t == 0) add(str, len);
        else if (t == 1) erase(str, len);
        else if (t == 2) output << search(str, len) << '\n';
        else if (t == 3) output << prefix(str, len) << '\n';
    }

    return 0;
}