Cod sursa(job #1647722)

Utilizator ELHoriaHoria Cretescu ELHoria Data 10 martie 2016 21:48:22
Problema Trie Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.3 kb
#include <fstream>
#include <map>
#include <algorithm>
#include <vector>
#include <string>
#include <iostream>
#include <memory>

using namespace std;

class Trie {

	struct TrieNode {
		map<char, unique_ptr<TrieNode> > next;
		int wordCount;
		int wordsInTree;
		TrieNode() {
			wordsInTree = 0;
			wordCount = 0;
		}
	};

	unique_ptr<TrieNode> root;

	void updateCount(const string& str, const int& value) {
		TrieNode *node = root.get();
		for (size_t i = 0; i < str.size(); i++) {
			node->wordsInTree += value;
			if (node->next.find(str[i]) == node->next.end()) {
				node->next[str[i]] = unique_ptr<TrieNode>(new TrieNode());
			}
			node = node->next[str[i]].get();
		}

		node->wordCount += value;
		node->wordsInTree += value;
	}

	void print(TrieNode* node, string str) const {
		if (node->wordCount > 0) {
			cout << str << "\n";
		}
		for (auto& w : node->next) {
			print(w.second.get(), str + w.first);
		}
	}

public:

	Trie() {
		root = unique_ptr<TrieNode>(new TrieNode());
	}

	void insert(const string& str) {
		updateCount(str, 1);
	}

	void remove(const string& str) {
		updateCount(str, -1);
	}

	int queryLcp(const string& str) const {
		TrieNode *node = root.get();
		int ans = 0;
		for (size_t i = 0; i < str.size(); i++) {
			if (node->next.find(str[i]) == node->next.end()) {
				break;
			}
			node = node->next[str[i]].get();
			if (node->wordsInTree > 0)
				ans = static_cast<int>(i + 1);
		}

		return ans;
	}

	int queryWordCount(const string& str) const {
		TrieNode *node = root.get();
		for (size_t i = 0; i < str.size(); i++) {
			auto it = node->next.find(str[i]);
			if (it == node->next.end()) {
				return 0;
			}
			node = node->next[str[i]].get();
		}
		return node->wordCount;
	}

	void printWords() const {
		print(root.get(), "");
		cout << "###\n";
	}

};

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