[자료구조] - 트라이(Trie)

 

목차
1. 들어가며
2. 트라이(Trie)란
3. 예제 적용

 

 

 

 

들어가며

 문자열 관련 문제에 단골로 등장하는 트라이(Trie)라는 자료구조에 대해 정리해보겠다. 트라이는 문자열 검색을 빠르게 실행할 수 있도록 설계한 트리 형태의 자료구조이다. 문자열이 여러개 주어지고 하나의 문자열이 다른 문자열에 속한다는 것을 확인하라거나, 여러개의 문자열을 사전에 저장한 후 검색같은 기능을 구현하라는 문제들이 모두 이 트라이 자료구조를 사용하는 문제라고 할 수 있다.

 

 

 

 

트라이(Trie)란?

 

 트라이는 트리의 일종으로 문자열의 검색에 특화된 자료구조이다. 예를들어 'apple', 'apply', 'car'의 3개의 단어를 트라이에 저장한다고 가정해보자.

 

 

 

 

 

 

 

 

 트라이는 root 노드아래에 단어들을 저장하는 방식이다. 'apple'와 'apply'는 'appl'까지 겹치는 글자가 있으므로 위와 같이 겹치는 글자까지는 동일한 노드를 지난다. 하지만 'car'의 경우 겹치는 부분도 없고 다른 겹치는 단어도 존재하지 않으므로 root아래에 혼자 존재하는 것을 확인할 수 있다.

 

 

 

 

 

 주의하여야 할 것은 트라이의 경우 트리이긴 하지만 이진 트리는 아니므로 3개 이상의 자식 노드도 가질 수 있다. 만약 주어지는 문자열이 알파벳으로 국한되는 경우 노드하나당 가질 수 있는 자식노드의 개수는 26개(대소문자 모두 포함시는 더 증가)가 된다고 볼 수 있다.

 

 

 

 

 그림을 통해 알아본 트라이는 위와 같다. 이제 코드로 트라이를 직접 구현해보자.

 

 

 

 

 

#include <iostream>
#include <unordered_map>
using namespace std;

class TrieNode {
public:
    unordered_map<char, TrieNode*> children;
    bool isEndOfWord;

    TrieNode() {
        isEndOfWord = false;
    }
};

class Trie {
private:
    TrieNode* root;

public:
    Trie() {
        root = new TrieNode();
    }

    void insert(string word) {
        TrieNode* curr = root;
        for (char ch : word) {
            if (curr->children.find(ch) == curr->children.end()) {
                curr->children[ch] = new TrieNode();
            }
            curr = curr->children[ch];
        }
        curr->isEndOfWord = true;
    }

    bool search(string word) {
        TrieNode* curr = root;
        for (char ch : word) {
            if (curr->children.find(ch) == curr->children.end()) {
                return false;
            }
            curr = curr->children[ch];
        }
        return curr != nullptr && curr->isEndOfWord;
    }
};

int main() {
    Trie trie;
    trie.insert("apple");
    cout << (trie.search("apple") ? "Found" : "Not found") << endl;   // 출력: Found
    cout << (trie.search("app") ? "Found" : "Not found") << endl;     // 출력: Not found
    cout << (trie.search("car") ? "Found" : "Not found") << endl; // 출력: Found
    return 0;
}

 

 

 

 

 글자하나를 나타내는 TrieNode에는 단어의 끝을 의미하는 boolean값과 자신의 자식 노드들을 가리키는 맵으로 구성되어 있다. 

 

 

 

 

 

 여기서 주의해야할 점은 단어의 끝을 의미하는 boolean 값이 true라고 해서 그것이 무조건 리프 노드(자식이 존재하지 않는 노드)를 의미하지 않는다는 것이다. 예를 들어 'apple'과 'app'이라는 단어들이 트라이 자료구조안에 들어가게 되면 두 단어는 'app'이라는 글자들이 겹치게 되고, 'app'의 두번째 'p'는 이 뒤로 'le'라는 자식 노드들이 더 존재함에도 불구하고 isEndOfWord라는 boolean 값이 true값으로 설정되게 된다.

 

 

 

 트라이를 사용해야하는 문제에서 위의 주의점을 꼭 인지하고 있어야 모든 테스트 케이스들을 통과할 수 있다. 이 점을 유의하며 예제 문제를 풀어보자.

 

 

 

 

예제 적용

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

5670번: 휴대폰 자판

 

 

 

 트라이의 구조를 이해하였다는 가정하에 풀 수 있는 예제를 풀어보자. 위 문제의 경우 주어진 단어들을 트라이에 입력하고 어떤 타이밍에 직접 자판을 눌러야하는지를 체크해줘야하는 문제이다. 바로 정답 코드를 보자.

 

 

 

#include <iostream>
#include <vector>
#include <unordered_map>
#include <iomanip>

using namespace std;

// 트라이를 활용한 문자열 검색 문제
// 기본적인 트라이를 구현 + 문제에서 원하는 일부 조건 추가


// 자식 노드가 아예 없는 경우 -> cnt 1증가 후 즉시 반환
// 자식 노드의 수가 1인 경우 -> cnt 변화 X 자동 진행
// 자식 노드의 수가 2이상인 경우 -> cnt 1증가 후 계속 진행
// 위의 3가지 경우의 수를 고려

// but, hell hello와 같이 이미 입력된 단어보다 더 긴 단어를 입력하고자 하는 경우가 문제
// 이경우 두번째 'l'에서 isEnd = true이고 자식의 개수가 1인경우에 수동으로 입력해줘야함
// 즉, 현재 노드가 isEnd = true이면서 자식의 개수가 1인 경우에 수동으로 입력하도록 설정

class TrieNode {
public:
	unordered_map<char, TrieNode*> children;
	bool isEnd;

	TrieNode() {
		isEnd = false;
	}
};

class Trie {
private:
	TrieNode* root;

public:
	Trie() {
		root = new TrieNode();
	}

	// TODO insert
	void insert(string word) {
		TrieNode* curNode = root; // 루트부터 돌면서 자리를 찾는다.
		for (char c : word) {
			if (curNode->children.find(c) == curNode->children.end()) { // 찾는 노드가 없는 경우
				curNode->children[c] = new TrieNode();
			}
			curNode = curNode->children[c]; // 다음 노드로 이동
		}
		curNode->isEnd = true; // 마지막 노드에서 리프노드임을 표시
	}

	int autoTyping(string word) { // 타이핑 횟수 반환
		TrieNode* curNode = root->children[word[0]];
		int cnt = 1; // 무조건 첫글자는 치므로 1부터 해야함
		for (int i = 1; i < word.length(); i++) {
			char c = word[i];
			if (curNode->children.size() >= 2) { // 자식노드의 수가 2이상인 경우 -> 이후의 경로가 여러가지인 경우 타이핑을 쳐야함
				++cnt;
			}
			else if (curNode->isEnd && curNode->children.size() == 1) { // 한 단어가 다른 단어에 속하는 경우, 경로가 한가지만 존재하더라도 수동으로 입력
				++cnt;
			}
			curNode = curNode->children[c];
		}

		return cnt; // 타이핑 횟수 반환
	}
};


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

	int n;
	while (cin >> n) {
		Trie *trie = new Trie();
		vector<string> arr;
		for (int i = 0; i < n; i++) {
			string str;
			cin >> str;
			arr.push_back(str);
			trie->insert(str);
		}
		

		double sum = 0;
		// 문자열 검색
		for (int i = 0; i < n; i++) {
			sum += trie->autoTyping(arr[i]);
		}

		cout << fixed << setprecision(2) << sum / n << "\n";


	}
}

 

 

 

 

 '자판을 어떤 때에 눌러야 하는가?'에 대해 초점을 맞추고 풀면 되는 문제이다. 문제의 조건에서 단어의 공통된 부분이 존재하면 그 글자까지 자동으로 입력된다고 하였으므로, 트라이의 구간이 나뉘는 부분에서 자판을 눌러야한다는 것을 알 수 있다. 트라이에서 구간이 나뉜다는 것은 해당 노드의 자식노드가 2개 이상 존재한다는 것으로도 생각할 수 있다.

 

 

 

 

 다만 앞서 언급하였듯이 'hello'와 'hell'같이 단어가 다른 단어에 속하는 경우, 두번째 'l'에서 구간이 나뉘지만 자식노드는 'o' 한개만 존재한다. 즉, 자식노드의 갯수는 1개지만 해당 노드의 isEndOfWord(boolean 값)이 true인 경우도 구간이 나뉘는 것으로 간주해줘야 한다는 것이다. 이 점을 유의하면 위와 같이 코드를 작성할 수 있다.