[백준/BOJ] 백준 27652번 : AB

https://www.acmicpc.net/problem/27652

27652번: AB

 

집합 A의 문자열을 관리하는 트라이와, 집합 B의 문자열을 관리하는 트라이를 사용했는데, 집합 B는 접미사를 확인하므로, 문자열 순서를 거꾸로 저장해서 관리했다. 그리고 find S 쿼리에 대해, 해당 문자열을 접두사 + 접미사로 나누어지는 경우를 모두 확인하는데, 접두사는 A 집합의 트라이에서, 접미사는 B 집합의 트라이에서 쿼리 한 값을 토대로 각 경우의 수를 구했다.

 

코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <utility>
#include <string>
using namespace std;

int q;

//트라이를 구현할때 구조체가 아닌 간단하게 트라이 구현

int root_a = 1; //그룹 a의 트라이 루트
int root_b = 2; //그룹 b의 트라이 루트
int node_number = 2;
int finish[1000 * 1000 + 5]; //[노드번호] = 해당 노드에서 단어가 끝나는 것 개수 (최대 노드 가능 개수 : 1000 * 1000)
int pass[1000 * 1000 + 5]; //[노드번호] = 해당 노드 지나는것 개수 (최대 노드 가능 개수 : 1000 * 1000)
int adj[1000 * 1000 + 5][26]; //[노드번호][알파벳 번호] = 해당 노드에서 해당 알파벳으로 연결되는 자식 노드 번호

vector<int> output;

void pre() {
	for (int i = 0; i < 1000 * 1000 + 5; i++) {
		finish[i] = 0;
		pass[i] = 0;

		for (int j = 0; j < 26; j++) {
			adj[i][j] = 0;
		}
	}
}

void trie_add(int root, string& input) {

	int here = root;

	for (int i = 0; i < input.size(); i++) {
		if (adj[here][input[i] - 'a'] == 0) { //트라이에 해당 자식노드가 없을때
			node_number++;
			adj[here][input[i] - 'a'] = node_number;
		}
		here = adj[here][input[i] - 'a'];
		pass[here]++; //해당 노드 지나는 개수 체크
	}

	finish[here]++; //단어가 끝나는 지점 체크
}

void trie_delete(int root, string& input) {

	int here = root;

	for (int i = 0; i < input.size(); i++) {
		here = adj[here][input[i] - 'a'];
		pass[here]--;
	}

	finish[here]--;
}

//root 트라이에서 input이랑 매치되는 정보 반환
vector<int> trie_find(int root, string& input) {
	int here = root;
	vector<int> ret(input.size(), 0); //[i] = 트라이에서 input[i]까지 매치되는것의 개수

	for (int i = 0; i < input.size(); i++) {
		int there = adj[here][input[i] - 'a'];
		if (there == 0 || pass[there] == 0) {
			break;
		}

		here = there;
		ret[i] = pass[here];
	}

	return ret;
}

int main()
{
	cin.tie(NULL);
	ios_base::sync_with_stdio(false);

	pre();

	cin >> q;

	for (int i = 0; i < q; i++) {
		string order;
		cin >> order;

		if (order == "add") {
			string group, input;
			cin >> group >> input;

			int root;
			if (group == "A") {
				root = root_a;
			}
			else {
				root = root_b;

				//B 집합의 경우 빠른 접미사 확인을 위해 거꾸로 해서 트라이에 넣는다
				reverse(input.begin(), input.end());
			}

			trie_add(root, input);
		}

		else if (order == "delete") {
			string group, input;
			cin >> group >> input;

			int root;
			if (group == "A") {
				root = root_a;
			}
			else {
				root = root_b;
				reverse(input.begin(), input.end());
			}

			trie_delete(root, input);
		}

		else {
			string input;
			cin >> input;

			//해당 단어가 트라이a와 트라이b에서 매치되는 정보를 이용해서 경우의 수를 구한다
			vector<int> info_a = trie_find(root_a, input);

			reverse(input.begin(), input.end());
			vector<int> info_b = trie_find(root_b, input);
			reverse(input.begin(), input.end());

			int result = 0;
			//접두사:0~j까지
			//접미사:j+1에서 input.size()-1 까지
			for (int j = 0; j < input.size() - 1; j++) {
				result += (info_a[j] * info_b[input.size() - 1 - 1 - j]);
			}
			output.push_back(result);
		}
	}

	for (int i = 0; i < output.size(); i++) {
		cout << output[i] << "\n";
	}

	return 0;
}