[Data Structure] 문자열을 위한 트리, 트라이(Trie)

· 서론

요즘 백준에서 문자열 문제를 주로 푸는데, 느끼는 것은 문자열은 알고리즘이 너무 많다는 것이다. 문자열 문제 중 골드쯤 된다 하면 대부분 문자열 특수 알고리즘을 사용하는 것이 대부분이다. 그걸 직접 생각해낼 수 있다면 이미 어나더레벨이 아닐까..

· 트라이(Trie) 자료구조

트라이는 문자열 관리에 유리한 트리 구조이며, 문자열 탐색에 빠른 성능을 보인다. 문자 또는 문자열에 순서가 있을 때 트라이 자료구조를 적용할 수 있으며, 간단하게 'apple'이라는 단어도 a-p-p-l-e라는 순서로 알파벳이 배치되기 때문에 트라이 자료구조를 적용할 수 있다.

위와 같이 영어 단어로 트라이 자료구조를 활용한다면 알파벳의 순서에 따라 트리에 원소를 삽입한다. 'ant', 'mad', 'made' 그리고 'mode'를 트라이에 저장한다면 다음과 같이 트리가 생성된다.

예시로는 이진 트리처럼 보이지만, 실제로는 이진 트리가 아니라는 것에 유의하자. 원하는 문자열을 탐색할 때 시작점인 Root node가 문자열의 첫 문자를 자식으로 가지고 있는지를 확인한다. 다음 문자를 값으로 하는 자식 노드가 존재한다면 해당 자식 노드가 문자열의 다음 문자를 자식으로 가지는지 계속 탐색하는 것이다. 그리고 찾을 문자열의 마지막 문자에 도달했을 때 해당 노드가 Last node인지 확인한다.

위 트리에서 m - a - d - e 순으로 나열된 부분 트리의 d 노드가 Last node라고 표시된 것을 확인할 수 있다. 이는 m - a - d를 마지막으로 하는 문자열이 현재 트리에 존재한다는 뜻이다. Last node는 해당 문자를 마지막으로 하는 문자열이 트리에 있다는 표시이다.

이를 구현한 코드 전문은 아래와 같다.

///////////////////////////////////////////////
//////////////////// Array ////////////////////
///////////////////////////////////////////////
#include <iostream>

class Node {
public:
	Node *child[26];	// 알파벳 26개
	bool last;

	Node() {
		for (int index = 0; index < 26; index++) {
			child[index] = NULL;
		}
		last = false;
	}
};

int n;

void insert(Node &node, std::string str, int index) {
	if (index == str.length()) {
		node.last = true;
		return;
	}

	int children = str[index] - 'a';

	if (node.child[children] == NULL) {
		node.child[children] = new Node();
	}

	insert(*node.child[children], str, index + 1);
}

bool find(Node& node, std::string str, int index) {
	if (index == str.length()) {
		return node.last;
	}
	if (node.child[str[index] - 'a'] == NULL) {
		return false;
	}
	return find(*node.child[str[index] - 'a'], str, index + 1);
}

void print_all(Node &node, std::string str, int index) {
	if (node.last) {
		std::cout << str << "\n";
	}

	for (int index = 0; index < 26; index++) {
		char ch;
		if (node.child[index] != NULL) {
			ch = index + 'a';

			print_all(*node.child[index], str + ch, index + 1);
		}
	}
}

int main() {
	std::ios::sync_with_stdio(false);
	std::cin.tie(NULL);
	std::cout.tie(NULL);

	std::cin >> n;

	Node root;

	std::string str;
	for (int count = 0; count < n; count++) {
		std::cin >> str;

		insert(root, str, 0);
	}

	std::string findStr;
	std::cin >> findStr;
	if (find(root, findStr, 0)) {
		std::cout << "YES\n";
	}
	else {
		std::cout << "NO\n";
	}

	print_all(root, "", 0);
}

· Map을 사용한 트라이(Trie)

위 Node 객체를 배열로 선언했을 때의 문제점은 26개와 같이 경우의 수가 적은 케이스에 적절하다는 것이고, 백준의 개미굴 문제와 같이 원소로 가질 수 있는 데이터의 경우의 수가 많은 경우에는 메모리 크기가 매우 커진다는 단점이 있다. 이때 사용할 수 있는 것이 Set/Map 자료구조이다.

이때 정렬이 필요하다면 Red-Black tree를 기반으로 하는 set/map STL을, 정렬이 필요하지 않다면 unordered_set/unordered_map STL을 사용하면 편리하다.

map STL을 활용하여 트라이 자료구조를 구현한 코드 전문은 아래와 같다.

///////////////////////////////////////////////
///////////////////// Map /////////////////////
///////////////////////////////////////////////
#include <iostream>
#include <map>

class Node {
public:
	std::map<char, Node> child;
	bool last;

	Node() {
		last = false;
	}
};

int n;

void insert(Node& node, std::string str, int index) {
	if (index == str.length()) {
		node.last = true;
		return;
	}

	if (node.child.count(str[index]) == 0) {
		node.child[str[index]] = Node();
	}

	insert(node.child[str[index]], str, index + 1);
}

bool find(Node& node, std::string str, int index) {
	if (index == str.length()) {
		return node.last;
	}

	if (!node.child.count(str[index])) {
		return false;
	}

	return find(node.child[str[index]], str, index + 1);
}
void print_all(Node& node, std::string str, int index) {
	if (node.last) {
		std::cout << str << "\n";
	}

	for (auto children : node.child) {
		print_all(children.second, str + children.first, index + 1);
	}
}
int main() {
	std::ios::sync_with_stdio(false);
	std::cin.tie(NULL);
	std::cout.tie(NULL);

	std::cin >> n;

	Node root;

	std::string str;
	for (int count = 0; count < n; count++) {
		std::cin >> str;

		insert(root, str, 0);
	}

	std::string findStr;
	std::cin >> findStr;
	if (find(root, findStr, 0)) {
		std::cout << "YES\n";
	}
	else {
		std::cout << "NO\n";
	}

	print_all(root, "", 0);
}