Cod sursa(job #1064556)

Utilizator cosmo0093Raduta Cosmin cosmo0093 Data 22 decembrie 2013 00:38:44
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.05 kb
#include <fstream>
#include <iostream>

class Trie {
	
	struct TrieNode {
		int sons, count;
		TrieNode* myArr[26];
		TrieNode() : sons(), count() {
			for(int i = 0; i < 26; i++) {
				myArr[i] = NULL;
			}
		}
	} *root;
	
	void _insert(std::string &szStr, unsigned i, TrieNode *c) {
		if(i == szStr.length()) {
			c->count++;
			return;
		}
		int aux = szStr[i] - 'a';
		if(c->myArr[aux] == NULL) {
			c->sons++;
			c->myArr[aux] = new TrieNode();	
		}
		_insert(szStr, i + 1, c->myArr[aux]);
	}
	
	int count(std::string &szStr, unsigned i, TrieNode *c) {
		if(i == szStr.length()) {
			return c->count;
		}
		int aux = szStr[i] - 'a';
		if(c->myArr[aux] == NULL) {
			return 0;
		}
		return count(szStr, i + 1, c->myArr[aux]);
	}
	
	int prefix(std::string &szStr, unsigned i, unsigned p, TrieNode *c) {
		if(i == szStr.length()) {
			return p;
		}
		int aux = szStr[i] - 'a';
		if(c->myArr[aux] == NULL) {
			return p;
		}
		return prefix(szStr, i + 1, p + 1, c->myArr[aux]);
	}
	
	bool _remove(std::string &szStr, unsigned i, TrieNode *c) {
		if(i == szStr.length()) {
			c->count--;
		} else if(_remove(szStr, i + 1, c->myArr[szStr[i] - 'a'])) {
			c->sons--;
			c->myArr[szStr[i] - 'a'] = NULL;
		}
		if(c->sons == 0 && c->count == 0 && c != root) {
			delete c;
			return true;
		}
		return false;
	}

public:
	
	Trie() : root(new TrieNode()) {
	}
	
	void insert(std::string &szStr) {
		_insert(szStr, 0, root);
	}
	
	int getCount(std::string &szStr) {
		return count(szStr, 0, root);
	}
	
	int getPrefix(std::string &szStr) {
		return prefix(szStr, 0, 0, root);
	}
	
	int remove(std::string &szStr) {
		_remove(szStr, 0, root);
	}
};



int main() {
	std::ifstream in("trie.in");
	std::ofstream out("trie.out");
	Trie a;
	int x; 
	std::string szStr;
	while(in >> x >> szStr) {
		if(x == 0) {
			a.insert(szStr);
		} else if(x == 1) {
			a.remove(szStr);
		} else if(x == 2) {
			out << a.getCount(szStr) << '\n';
		} else {
			out << a.getPrefix(szStr) << '\n';
		}
	}
	return 0;
}