Cod sursa(job #2662678)

Utilizator betybety bety bety Data 24 octombrie 2020 12:31:24
Problema Trie Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.24 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;

}