Cod sursa(job #1237026)

Utilizator ELHoriaHoria Cretescu ELHoria Data 2 octombrie 2014 23:45:34
Problema Trie Scor 0
Compilator cpp Status done
Runda Arhiva educationala Marime 1.82 kb
#include <fstream>
#include <map>
#include <vector>
#include <string>

using namespace std;

class Trie {

	struct TrieNode {
		map<char, TrieNode*> next;
		int wordCount;
		
		TrieNode() {
			wordCount = 0;
		}
	};
	
	TrieNode* root;
	
	void updateCount(const string& str, const int& value) {
		TrieNode *node = root;
		for (size_t i = 2; i < str.size(); i++) {
			auto it = node->next.find(str[i]);
			if (it == node->next.end()) {
				node->next[str[i]] = new TrieNode();
			}
			node = node->next[str[i]];
		}
		
		node->wordCount += val;
	}
	
	
	public:
		
		Trie() {
			root = new TrieNode();
		}
		
		void insert(const string& str) {
			updateCount(str, 2, 1);
		}
		
		void remove(const string& str) {
			updateCount(str, 2, -1);
		}
		
		int queryLcp(const string& str) {
			TrieNode *node = root;
			int ans = 0;
			for (size_t i = 2; i < str.size(); i++) {
				auto it = node->next.find(str[i]);
				if (it == node->next.end()) {
					break;
				}
				node = node->next[str[i]];
				if (node->wordCount > 0) {
					ans = max(ans,static_cast<int>(i - 1));
				}
			}
			
			return ans;
		}
		
		int queryWordCount(const string& str) {
			TrieNode *node = root;
			for (size_t i = 2; i < str.size(); i++) {
				auto it = node->next.find(str[i]);
				if (it == node->next.end()) {
					return 0;
				}
				node = node->next[str[i]];
			}
			return node->wordCount;
		}
		
};

int main() {
	ifstream cin("trie.in");
	ofstream cout("trie.out");
	Trie trie;
	string str;
	while (getline(cin, str)) {
		if (str[0] == '0') {
			trie.insert(str);
		} else
		if (str[0] == '1') {
			trie.remove(str);
		} else
		if (str[0] == '2') {
			cout << trie.queryWordCount(str) << "\n";
		} else
		if (str[0] == '3') {
			cout << trie.queryLcp(str) << "\n";
		}
	}
	return 0;
}