본문 바로가기

Golang/Golang 자료구조

[Golang 자료구조] 트라이(Trie) 구현하기

728x90

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에는 없지만 멋있는 자료구조를 구현해보도록 하자.

728x90