Cod sursa(job #3032007)

Utilizator livliviLivia Magureanu livlivi Data 21 martie 2023 11:42:06
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.95 kb
#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

struct Node {
	int nr_cuv = 0;
	int nr_ap = 0;
	vector<Node*> fii;

	Node() : fii(26) {}
};

class Trie {
public:
	void insert(string s) { root_ = insert_(root_, s.c_str()); }
	void erase(string s) { root_ = erase_(root_, s.c_str()); }
	int query_prefix(string s) { return query_prefix_(root_, s.c_str()); }
	int query_word(string s) { return query_word_(root_, s.c_str()); }

private:
	Node* root_ = nullptr;

	Node* insert_(Node* node, const char* s);
	Node* erase_(Node* node, const char* s);
	int query_word_(Node* node, const char* s);
	int query_prefix_(Node* node, const char* s);
};

Node* Trie::insert_(Node* node, const char* s) {
	if (node == nullptr) { node = new Node; }
	
	node->nr_ap++;

	if (s[0] == '\0') {
		node->nr_cuv++;
	} else {
		node->fii[s[0] - 'a'] = insert_(node->fii[s[0] - 'a'], s + 1);
	}

	return node;
}

Node* Trie::erase_(Node* node, const char* s) {
	if (node == nullptr) { return node; }

	if (s[0] == '\0') {
		node->nr_cuv--;
	} else {
		node->fii[s[0] - 'a'] = erase_(node->fii[s[0] - 'a'], s + 1);
	}

	node->nr_ap--;
	if (node->nr_ap == 0) {
		delete node;
		node = nullptr;
	}

	return node;
}

int Trie::query_word_(Node* node, const char* s) {
	if (node == nullptr) { return 0; }

	if (s[0] == '\0') {
		return node->nr_cuv;
	} else {
		return query_word_(node->fii[s[0] - 'a'], s + 1);
	}
}

int Trie::query_prefix_(Node* node, const char* s) {
	if (node == nullptr || s[0] == '\0') { return 0; }

	if (node->fii[s[0] - 'a'] == nullptr) {
		return 0;
	} else {
		return 1 + query_prefix_(node->fii[s[0] - 'a'], s + 1);
	}
}

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

	Trie trie;

	int op;
	while (in >> op) {
		string s; in >> s;

		if (op == 0) {
			trie.insert(s);
		} else if (op == 1) {
			trie.erase(s);
		} else if (op == 2) {
			out << trie.query_word(s) << "\n";
		} else {
			out << trie.query_prefix(s) << "\n";
		}
	}

	return 0;
}