문자열 관리 자료구조 - 트라이(Trie) 개념/구현

 C++로 PS를 공부할 때나 코딩테스트를 볼 때 문자열 문제가 나오면 일단 짜증이난다.

파이썬은 문자열 다루기가 편하다는데 C++는 아무래도 문자열을 다룰때 힘든 면이 많다. 문자열 문제가 나오면 파이썬을 쓰고 다른 문제에선 C++을 쓰는 사람도 있을 정도다.

아무튼, 트라이는 문자열을 다루는 방식 중에서도 특히 문자검색을 할 때 가장 유용하고 효과적인 자료구조가 바로 트라이 이다. 

 

먼저 트라이는 문자열로 트리를 만든다고 보면 된다. AB, ABC, BCD, ABD 라는 단어가 들어온다고 치자. 그렇다면

 

       root

       /    \

      A     B

     /         \

   B*          C

  /   \           \

C *  D*       D*           

 

이런식으로 트리를 구성한다. 이렇게 하고 단어가 끝나는 글자들에는 표시를 해둔다. 즉 , 단어를 찾을 때 표시가 되어있는 글자라면 끝났음을 알 수 있는 것이다.

코드는 다음과 같다.

 

#include <iostream>

using namespace std;

struct TRIE {
	bool finish;
	TRIE* node[26];

	TRIE() {
		finish = false;
		for (int i = 0; i < 26; i++)
			node[i] = NULL;
	}

	void insert(string str,int idx) {
		
		if (idx == str.size() - 1) {
			finish = true;
		}
		if(idx == str.size())
			return;
		
		int index = str[idx] - 'A';
		
		if (node[index] == NULL) {
			node[index] = new TRIE();
		}
		node[index]->insert(str, idx + 1);
	}

	bool find(string str, int idx) {
		if (idx == str.size() - 1) {
			if (finish == true) return true;
			else return false;
		}
		int index = str[idx] - 'A';

		if (node[index] == NULL) return false;
		return node[index]->find(str, idx + 1);
	}
};

int main() {
	TRIE* root = new TRIE();
	root->insert("ABC", 0);
	root->insert("CDE", 0);
	root->insert("EFG", 0);
	cout << root->find("ABC", 0) << endl;
	cout << root->find("CDE", 0) << endl;
	cout << root->find("EFG", 0) << endl;
	cout << root->find("ABCD", 0) << endl;

	return 0;
}