Cod sursa(job #2654724)

Utilizator raikadoCri Lu raikado Data 2 octombrie 2020 02:52:55
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 1.9 kb
#include <fstream>
#include <iostream>
#include <string>
#include <vector>
#include <array>
#include <algorithm>

using namespace std;

static const int NCHARS = 26;
static const char FIRSTCHAR = 'a';


struct Node
{
	array<Node*, NCHARS> desc{nullptr};
	int count = 0;
};

class Trie
{
	Node root;

public:
	void insert(const string &w)
	{
		Node* node = &root;
		for (size_t i = 0; i < w.size(); i++)
		{
			int idx = w[i] - FIRSTCHAR;
			if (node->desc[idx] == nullptr)
				node->desc[idx] = new Node;
			node = node->desc[idx];
		}
		node->count++;
	}

	void erase(const string &w)
	{
		Node* node = &root;
		vector<Node*> path; path.reserve(w.size());
		path.push_back(node);

		for (size_t i = 0; i < w.size(); i++)
		{
			int idx = w[i] - FIRSTCHAR;
			node = node->desc[idx];
			path.push_back(node);
		}
		node->count--;

		size_t i = w.size();
		do {
			i--;
			if (path[i+1]->count == 0 && 
				all_of(path[i+1]->desc.begin(), path[i+1]->desc.end(), [](Node* n){return n == nullptr;}))
			{
				delete path[i+1];
				path[i]->desc[w[i] - FIRSTCHAR] = nullptr;
			} else
				break;
		} while (i > 0);

	}

	int count(const string &w)
	{
		Node* node = &root;
		for (size_t i = 0; i < w.size(); i++)
		{
			int idx = w[i] - FIRSTCHAR;
			if (node->desc[idx] == nullptr)
				return 0;
			node = node->desc[idx];
		}
		return node->count;
	}

	int prefix(const string &w)
	{
		int cnt = 0;

		Node* node = &root;
		for (size_t i = 0; i < w.size(); i++)
		{
			int idx = w[i] - FIRSTCHAR;
			if (node->desc[idx] == nullptr)
				return cnt;
			node = node->desc[idx];
			cnt++;
		}
		return cnt;
	}
};

int main(int argc, char const *argv[])
{
	ifstream fin("trie.in");
	ofstream fout("trie.out");

	int op;
	string w;
	Trie T;
	while (fin >> op >> w)
	{
		if (op == 0)
			T.insert(w);
		else if (op == 1)
			T.erase(w);
		else if (op == 2)
			fout << T.count(w) << '\n';
		else if (op == 3)
			fout << T.prefix(w) << '\n';
	}

	return 0;
}