Golang은 내장 자료구조가 별로 없다. 하다못해 스택도 없다.
그래서 golang으로 몇몇 자료구조를 직접 구현하여 사용해보는 시리즈를 진행하고자 한다.
그 첫 번째는 간단하게 구현 및 테스트를 할 수 있는 트라이(Trie)이다.
Trie
트라이를 간략하게 설명하면 여러 단어들의 집합을 하나의 트리로 관리하여 어떤 단어 $S$의 존재 여부를 $O(|S|)$의 시간 복잡도로 찾을 수 있는 자료구조이다.
비교적 간단하게 구현할 수 있고 트라이를 이용한 다양한 알고리즘 테크닉을 쓸 수가 있어 PS에서 인기가 많은 편이다.
이 포스트에서는 PS 수준에서 활용할 수 있는 간단한 트라이를 구현할 것이다.
Node 구조체
트라이에 필요한 노드 구조체를 정의해보자.
type trieNode struct {
childs [26]*trieNode
}
알파벳 단어로 한정시킬 것이므로 자식으로는 26개의 알파벳이 올 수 있다. 그래서 26개의 자식 노드를 준비한다.
트라이의 용도에 따라 멤버가 더 필요할 수 있지만 지금은 이것으로도 충분하다.
삽입 함수 구현
단어들을 넣은 트리를 구성해야 하기 때문에 문자열을 삽입하는 함수가 필요하다.
func insert(node *trieNode, text *string, idx int) {
length := len(*text) // 문자열 길이
if length <= idx {
// 문자열 끝에 도달했으면 재귀 종료
return
}
// 현재 문자가 몇번째 순서인지 구한다.
cur := (*text)[idx] - 'a'
if node.childs[cur] == nil {
// cur에 해당하는 자식 노드가 없으면 새로 만든다.
// 이미 있는 경우 현재까지의 prefix가 있다는 뜻이므로 만들지 않아도 된다.
node.childs[cur] = &trieNode{}
}
// 다음 자식 노드로 재귀 진행
insert(node.childs[cur], text, idx+1)
}
현재 노드 node, 삽입 대상 문자열인 text, 삽입 중인 인덱스 idx를 매개변수로 넣는다.
문자열 끝에 도달했는지 검사하여 끝에 도달하면 종료한다.
끝이 아니라면 삽입 대상 알파벳을 [0, 25] 범위로 변환한다.
그리고 현재 알파벳에 해당하는 자식 노드가 존재하는지 살펴보고, 없으면 생성한 후 그 노드로 재귀를 진행한다.
접두사 매칭 함수 구현
이 포스트에서는 트라이에 특정 접두사가 존재하는지 판별하는 함수를 구현할 것이다.
func search_pref(node *trieNode, target *string, idx int) bool {
if len(*target) <= idx {
// 문자열 끝에 도달했으면 접두사라는 뜻이므로 true 반환
return true
}
nxt := (*target)[idx] - 'a'
if node.childs[nxt] == nil {
// 자식 노드가 없다는 것은 길이 막힘. 즉, target을 접두사로 가지는 문자열이 트라이에 없다.
return false
}
return search_pref(node.childs[nxt], target, idx+1)
}
접두사로 검사할 target이 트라이 끝으로 갈 수 있는지 검사하도록 하여 도달 여부를 반환한다.
테스트 - 백준 1446 접두사 찾기
이 트라이가 제대로 구현한 것인지 백준에서 채점을 받아보도록 하자.
백준 온라인 저지 정도면 충분히 강한 데이터가 있다고 믿을 수 있기 때문에 백준에서의 채점을 통과하면 만든 자료구조는 충분히 신뢰할 수 있다고 볼 수 있다.
https://www.acmicpc.net/problem/14426
위 문제는 $N$개의 단어로 구성된 집합을 대상으로 $M$개의 단어 중 몇 개가 접두사가 될 수 있는지 갯수를 세는 문제다.
아래와 같이 코드를 작성했다.
이정도면 꽤 나쁘지 않은 속도이다.
이렇게 간단하게 트라이를 구현해보고 백준을 통해 검증했다.
다음에도 golang에는 없지만 멋있는 자료구조를 구현해보도록 하자.