Cod sursa(job #2137814)

Utilizator rares96cheseliRares Cheseli rares96cheseli Data 21 februarie 2018 03:35:14
Problema Trie Scor 80
Compilator cpp Status done
Runda Arhiva educationala Marime 2.37 kb
#include <bits/stdc++.h>

#define LMAX 26

using namespace std;

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

struct TrieNode {
    /* number of words that pass through this node */
    int wordCount;
    TrieNode* child[LMAX];

    TrieNode(int val) {
        wordCount = val;
        for (int i = 0; i < LMAX; ++i) {
            child[i] = nullptr;
        }
    }
};

TrieNode* root = new TrieNode(0);

void insertWord(TrieNode *node, const char* word) {
    if (word[0] != NULL) {
        int c = word[0] - 'a';

        if (node->child[c] == nullptr) {
            node->child[c] = new TrieNode(1);
        } else {
            node->child[c]->wordCount++;
        }

        insertWord(node->child[c], word + 1);
    }
}

bool eraseWord(TrieNode* node, const char* word) {
    if (word[0] != NULL) {
        int c = word[0] - 'a';
        eraseWord(node->child[c], word + 1);

        if (node->child[c]->wordCount == 1) {
            delete node->child[c];
            node->child[c] = nullptr;
        } else {
            node->child[c]->wordCount--;
        }
    }
}

int wordCnt(TrieNode* node, const char* word) {
    int c = word[0] - 'a';
    // TrieNode* next = ;

    if (node->child[c] == nullptr) {
        return 0;
    }

    if (word[1] == NULL) {
        int sum = 0;

        /* count the words that pass through the current node and are larger that "word" */
        for (int i = 0; i < LMAX; ++i) {
            if (node->child[c]->child[i]) {
                sum += node->child[c]->child[i]->wordCount;
            }
        }
        return node->child[c]->wordCount - sum;
    }
    return wordCnt(node->child[c], word + 1);
}

int findPrefix(TrieNode* node, const char* word) {
    if (word[0] == NULL) {
        return 0;
    }

    int c = word[0] - 'a';
    if (node->child[c] == nullptr) {
        return 0;
    }

    return 1 + findPrefix(node->child[c], word + 1);
}

int main() {
    int type;
    char word[21];

    while (fin >> type >> word) {
        if (type == 0) {
            insertWord(root, word);
        } else if (type == 1) {
            eraseWord(root, word);
        } else if (type == 2) {
            fout << wordCnt(root, word) << '\n';
        } else if (type == 3) {
            fout << findPrefix(root, word) << '\n';
        }
    }
    return 0;
}