Cod sursa(job #1786787)

Utilizator whoasdas dasdas who Data 23 octombrie 2016 17:58:36
Problema Trie Scor 5
Compilator cpp Status done
Runda Arhiva educationala Marime 2.45 kb
#define IA_PROB "trie"

#include <bits/stdc++.h>

using namespace std;

int num_nodes;

class trie {
public:
	void insert(const string &s);
	void remove(const string &s);
	int count(const string &s);
	int prefix(const string &s);
private:
	static bool validate_char(char c) { return c >= 'a' && c <= 'z'; }
	static bool validate_str(const string &s) { return all_of(s.begin(), s.end(), validate_char); }
	static int c2i(char c) { return  c - 'a';}

	typedef string::const_iterator str_cit;

	class node {
	public:
		int num_exact, num_prefix;
		node* get_next(char c) {
			return next_arr[c2i(c)].get();
		}
		void make_next(char c) {
			assert(next_arr[c2i(c)] == nullptr);
			next_arr[c2i(c)] = unique_ptr<node> (new node);
		}
		void delete_next(char c) {
			next_arr[c2i(c)] = nullptr;
		}

		node(): num_exact(0), num_prefix(0) {num_nodes++;};
		~node() {num_nodes--;}
	private:
		array<unique_ptr<node>, 32> next_arr;
	} root;

	pair<node*, str_cit> find(const string &s);
};

void trie::insert(const string &s) {
	node *cur = &root;
	for (auto it = s.begin(); it != s.end(); it++) {
		assert(validate_char(*it));
		cur->num_prefix++;

		if (!cur->get_next(*it)) {
			cur->make_next(*it);
		}
		cur = cur->get_next(*it);
	}
	cur->num_prefix++;
	cur->num_exact++;
}

pair<trie::node*, trie::str_cit> trie::find(const string &s) {
	node *cur = &root;
	str_cit it;
	for (it = s.begin(); cur && it != s.end(); it++) {
		cur = cur->get_next(*it);
	}
	return make_pair(cur, it-1);
}


int trie::count(const string &s) {
	auto ret = find(s);
	return ret.first ? ret.first->num_exact : 0;
}

int trie::prefix(const string &s) {
	auto ret = find(s);
	return ret.second - s.begin();
}

void trie::remove(const string &s) {
	if (count(s) == 0) {
		return;
	}
	node *cur = &root;
	for (auto it = s.begin(); cur && it != s.end(); it++) {
		cur->num_prefix--;
		assert(cur->num_prefix > 0 || (cur == &root && cur->num_prefix == 0));
		node *next = cur->get_next(*it);
		if (next->num_prefix == 1) {
			cur->delete_next(*it);
		}
		cur = cur->get_next(*it);
	}
	if (cur) {
		cur->num_prefix--;
		cur->num_exact--;
		assert(cur->num_prefix > 0);
		assert(cur->num_exact >= 0);
	}
}



int main() {
	ifstream in(IA_PROB".in");
	ofstream out(IA_PROB".out");

	int op;
	string s;

	{
	trie t;

	while (in>>op>>s) {
		switch (op) {
		case 0: t.insert(s); break;
		case 1: t.remove(s); break;
		case 2: out<<t.count(s)<<"\n"; break;
		case 3: out<<t.prefix(s)<<"\n"; break;
		}
	}
	}
	cerr<<num_nodes<<endl;

	return 0;
}