Cod sursa(job #1673160)

Utilizator Marius_mFMI-M2 Marius Melemciuc Marius_m Data 3 aprilie 2016 15:05:57
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 4.63 kb
/*
 * =====================================================================================
 *
 *       Filename:  IA_trie.cpp
 *
 *    Description:  http://www.infoarena.ro/problema/trie
 *
 *        Version:  1.0
 *        Created:  03/24/2016 10:50:30 AM
 *       Revision:  none
 *       Compiler:  gcc/g++
 *
 *         Author:  Marius-Constantin Melemciuc 
 *          email: 
 *   Organization: 
 *
 * =====================================================================================
 */
  
#include <iostream>
#include <string>
#include <map>
#include <fstream>
  
using namespace std;

// This verison of TrieNode holds a map instead of a vector => 
// is more space efficient. 

struct TrieNode {
    int countWords;
    map<char, TrieNode*> child; // is implemented as a BST,  
                                // so it has O(logN) for push, pop, find     
                                // I don't use hash - unordered_map, because it gives TLE    
    TrieNode() {
        countWords = 0;
    }
};
  
void pushToTrie(TrieNode* root, int letter, string& str) {
    if (root == NULL) {
        return;
    }
    if (letter == str.length()) { // finished the current word
        root->countWords++;
        return;
    }
    if (root->child.find(str[letter]) == root->child.end()) { // if there isn't a child with character str[letter]
        root->child[str[letter]] = new TrieNode;              // create a child with that character
    }
    pushToTrie(root->child[str[letter]], letter + 1, str);    // continue the recursion with the next character from word
}
  
bool popFromTrie(TrieNode* root, TrieNode* rootNode, int letter, string& str) {
    if (root == NULL) {
        return false;
    }
    if (letter == str.length()) { // finished the current word
        root->countWords--;
    } else {
        if (root->child.find(str[letter]) != root->child.end() &&  // if we find that character in map
            popFromTrie(root->child[str[letter]],
                        rootNode,
                        letter + 1,
                        str) == true) {  // if I eliminated the child
            root->child.erase(root->child.find(str[letter])); // erase the child
        }
    }
  
    if (root->child.size() == 0 && // if the node doesn't have any childs and 
        root->countWords == 0   && // doesn't have words that end in him and 
        root != rootNode) {        // is not the trie root
  
        delete root;               // delete the node
        return true;
    }
    return false;
}
  
int getNumberOfWords(TrieNode* root, int letter, string& str) {
    if (root == NULL) {
        return 0;
    }
    if (letter == str.length()) { // we have reached the end of the word
        return root->countWords;  // return the number of words that terminates in this node
    }
    if (root->child.find(str[letter]) == root->child.end()) { // if there isn't a child with the current character - str[letter]
        return 0; // we can't go further for searching the words' next character
    }
  
    return getNumberOfWords(root->child[str[letter]], 
                            letter + 1, // continue with the next character
                            str);
}
  
int getLengthOfPrefix(TrieNode* root, int letter, string& str) {
    if (root == NULL) {
        return 0;
    }
    if (letter == str.length() || root->child.find(str[letter]) == root->child.end()) { // if I've finished the current word or
        return letter;            // there isn't a child with words' current character, return the length found until now
    }
  
    return getLengthOfPrefix(root->child[str[letter]], letter + 1, str); // go to the child with the next character
}
  
/*
bool isInTrie(TrieNode* root, string str) {
    if (root == NULL) {
        return false;
    }
    if (str.length() < 1) {
        cout << root->countWords << " ";
        return true;
    }
    if (root->child.find(str[0]) == root->child.end()) {
        return false;
    }
    return isInTrie(root->child[str[0]], str.substr(1, str.length()));
}
*/
  
int main() {
    int operation;
    string str;
    ifstream input("trie.in");
    ofstream output("trie.out");
  
    TrieNode* rootNode;
    rootNode = new TrieNode;
  
    while (input >> operation >> str) {
        if (operation == 0) {
            pushToTrie(rootNode, 0, str);
        }
        if (operation == 1) {
            popFromTrie(rootNode, rootNode, 0, str);
        }
        if (operation == 2) {
            output << getNumberOfWords(rootNode, 0, str) << "\n";
        }
        if (operation == 3) {
            output << getLengthOfPrefix(rootNode, 0, str) << "\n";
        }
    } /* while */
  
    input.close();
    output.close();
 
    return 0;
}