Trie

트리(Trie) 자료구조

이 글에서는 Tree와의 구분을 위해 Trie라는 영문 표기를 사용합니다. 이 둘은 똑같이 트리1라 발음하며, 오타가 아닌 서로 다른 자료 구조입니다.

Trie는 특히 영단어들을 저장하는데 유용한 자료구조로 다음과 같은 장점을 가진다.

  • 값을 찾는 것은 최악의 경우에서 더 나은 시간 복잡도를 보인다.
  • 해시테이블과는 달리 키 충돌을 염려하지 않아도 된다.
  • 특정한 요소에 다다르는 경로를 찾기위한 부가적인 해시 알고리듬을 요구하지 않는다.
  • 트리 자료구조는 알파벳순으로 정렬하기가 매우 쉽다.

일반적인 Tree들과 마찬가지로 Trie는 노드들로 구성된다. 이 글에서 각 노드는 TrieNode로, 전체 트리는 Trie로 구현할 것이다. 예를 들어 cute라는 단어는 다음과 같은 일련의 노드로 표현된다.

c -> u -> t -> e

노드

먼저 기본적인 TrieNode 클래스의 구성은 다음과 같다. 여기서 T는 사전의 키로 사용될 것이기 때문에 Hashable을 준수해야 한다.

class TrieNode<T: Hashable> {
  var value: T?
  weak var parent: TrieNode? 
  var children: [T: TrieNode] = [:]

  init(value: T?, parent: TrieNode? = nil) {
    self.value = value
    self.parent = parent
  }

  // 위에서 커스텀 designated initializer를 제공했으므로, 
  // init()은 편의 이니셜라이저로 따로 구현해야 한다. 
  convenience init() {
    self.init(value: nil, parent: nil)
  }
}

노드에 자식을 더하기 위해서는 보통의 Tree가 하는 것처럼 다음과 같은 메소드를 추가할 수 있다. 물론 몇 가지 다른 점이 있어야하겠지만, 이는 노드의 기능보다는, 트리의 기능으로 구현할 것이다.

extension TrieNode {
  func add(_ child: T) {
    guard children[child] == nil else { return }
    children[child] = TrieNode(value: child, parent: self)
  }
}

Trie

Trie 구조는 단순하다. 값을 가지지 않는 루트 노드를 가지는 트리구조이다.

class Trie {
  typealias Node = TrieNode<Character>
  fileprivate let root = Node()
}

트리에 새 노드를 추가해보자. 주어진 단어를 한글자씩 분해하면서 트리 구조를 따라 추가해준다.

extension Trie {
  func insert(word: String) {
    guard !word.characters.isEmpty else { return }
    var currentNode: Node? = root
    let characters = Array(word.lowercased().characters)
    var currentIndex = 0

    while currentIndex < characters.count {
      let char = characters[currentIndex]
      if currentNode?.children[char] == nil {
        currentNode.add(char)
      }
      currentNode = currentNode?.children[char]
      currentIndex += 1
    }
    currentNode?.isTerminating = true
  }
}

여기서 주목할 점은 최종단계에서 isTerminating값을 설정해준다는 것이다. 이는 hot, hottest가 있는 경우, 이 두 단어는 같은 경로상에 위치하게 되는데, 이 때 ho, hotte 등은 포함된 문자열이 아닌 것으로 판단하기 위해서 탐색의 끝을 매치하기 위해 필요하다.

탐색

이번에는 특정 단어가 트리 내에 있는지를 탐색한다. 트리에서는 이 작업이 최악의 경우에도 단어의 길이에 해당하는 값을 얻는다. 삽입과 마찬가지로 주어진 단어를 한 글자씩 분해하여 트리를 따라 내려간다. 중간에 찾을 수 없는 부분열을 만나면 즉시 false를 리턴할 수 있다. 주어진 단어의 깊이 만큼 진행했다 하더라도, 해당 노드의 isTerminating 값이 false라면 경로 상에 위치는 하지만 입력된 적은 없는 단어임을 체크해야 한다.

extension Trie {
  func contains(word: String) -> Bool {
    guard !word.characters.isEmplty else { return false }
    var currentNode: Node? = root
    var currentIndex = 0
    let characters = Array(word.lowercased().characters)
    while currentIndex < characters.count {
      let char = characters[currentIndex]
      if currentNode?.children[char] == nil { return false }
      currentNode = currentNode?.children[char]
      currentIndex += 1
    }
    if (currentNode?.isTerminating ?? false) { return true}
    return false
  }
}

테스트를 해보자.

let trie = Trie()
["cure", "care", "cute", "cut", "abort"].forEach{ trie.insert(word: $0) }

["cure", "curse", "cat", "cute", "about"].forEach{ word in 
  let i = trie.contains(word: word) ? "" : "not "
  print(" \(word): \(i)found ")
}
/*
cure: found
curse: not found
cat: not found
cute: found
about: not found
*/

순회

Swift 3.1에서는 제네릭 타입에 대해서 특정 조건에서만 확장을 적용할 수 있지만, 현재 사용가능한 플랫폼이 3.0.2 인 관계로 Trie 클래스를 확장하여 순회로직을 작성한다.

알파벳순으로 순회하는 메소드를 작성해보자. 각 단계에서 자식의 키들을 정렬하고, 정렬된 각 키들을 순서대로 모아나가면서 isTerminating으로 단어의 완성을 체크하며 출력해나간다.

extension Trie {
  func traverse(_ f: @escaping ([Character]) -> Void) {
    func visit(from n: Node?, _ p: [Character?]) {
      guard let node = n else { return }
      let prefixes = p + [node.value]
      if node.isTerminating {
        f(prefixes.flatMap{$0})
      }
      for key in node.children.keys.sorted() {
        visit(from: node.children[key]!, prefixes)
      }
    }
    visit(from: root, [])
  }
}

테스트:

trie.traverse{
  print(String($0))
}
/*
abort
care
cure
cut
cute
*/

도전 과제

  1. 같은 문자열이 2개 이상 들어가는 경우, 이를 카운트할 수 있는 구조를 만들어 보자.
  2. 삭제 기능을 추가해보자.
  3. Trie 간의 결합, 빼기, 공통집합 구하기 기능을 구현해보자.

  1. 하지만 많은 사람들이 “트라이”라고 발음하고 있다는….