Trie 자료구조 – Swift

for c in currentNode.children.keys.sorted() {트리(trie)자료 구조는 트리(tree)와 발음이 비슷한데, 실제로도 tree와 비슷한 형태로 문자열이나 이와 유사한 형태의 연속열들을 저장할 수 있는 자료 구조이다. 보통 영단어들을 저장할 때 유용한데, 다음과 같은 장점들을 가진다.

  • 특정한 값을 찾을 때, 다른 구조들에 비해서 최악의 경우에서도 대개 더 나은 시간 복잡도를 보인다.
  • 해시테이블과 달리 키 충돌을 고려할 필요가 없다.
  • 특정한 요소에 다다르는 경로를 찾기위한 부가적인 해시 알고리듬을 요구하지 않는다.
  • 알파벳순으로 정렬하기가 매우 쉽다.

Trie는 어떤 단어의 첫글자와 그 다음에 올 수 있는 글자, 그리고 그 다음글자… 이런 식으로 각 글자들을 순서대로 직렬로 연결한다. 그리고 특정한 글자에서 다음글자는 0개 이상일 수 있다.

트리(trie)는 tree와 비슷하게 노드들의 연결로 구성된다. 영단어를 저장하는 trie의 경우 각각의 노드는 낱개의 알파벳문자 값을 가지며, 다음 글자로 오게될 문자들을 참조한다. 따라서 구현의 시작은 트리노드가 되어야 한다.

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

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

  convenience init() {
    self.init(value:nil, parent:nil)
  }
}

그리고 trie 자체는 nil value를 가지는 빈 루트 노드를 갖는 형식으로 초기화될 수 있다.

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

단어 삽입

트리에 단어를 하나 추가한다고 생각해보자. 단어의 첫 글자는 루트의 자식이 되고, 그 다음 글자는 그 글자의 자식이되고… 이런 식으로 일련의 체인이 만들어질 것이다. 그런데 app과 apple 처럼 앞 부분을 공유하는 단어들이 같은 트리에 들어가는 경우에는 이미 존재하는 노드까지는 중복으로 설정할 필요가 없겠고, 다음 노드로 진행하면 된다. 전반적인 처리 방식은 연결 리스트(linked list)와 유사할 것이다.

extension TrieNode {
  func add(_ value: T) -> TrieNode {
    // 현재 노드에서 글자 하나를 추가한다. 
    // 리턴되는 노드는 해당 글자값을 갖는 서브 노드
    // 이미 그 글자를 가지고 있다면 해당 서브노드를 바로 리턴한다.
    guard children[value] == nil
      else { return children[value]! }
    let newNode = TrieNode(value: value, parent: self)
    children[value] = newNode
    return newNode
  }
}

extension Trie {
  func insert(word: String) {
    // 트리에 새 단어를 추가한다.
    // 단어의 앞 글자부터 루트에 추가해나가면서
    // 단어의 끝까지 글자를 추가한 후 마지막 노드에 
    // 종결 플래그를 설정한다. 
    var currentNode = root
    var currentIndex = word.startIndex
    while currentIndex < word.endIndex {
      let c = word[currentIndex]
      let nextNode = currentNode.add(c)
      currentIndex = word.index(after:currentIndex)
    }
    currentIndex.isTerminating = true
  }

단어 검사

다음은 trie에서 특정 단어를 포함하고 있는지 여부를 판단하는 메소드를 작성해보자. 삽입때와 마찬가지로 루트로부터 각 글자에 해당하는 자식이 있는지를 계속 검사해나간다. 마지막 글자까지 포함하고 있다면, 마지막 글자를 포함하는 해당 노드가 종료노드인지도 검사해야 한다. “hottest”를 삽입한 후에 “hot”을 검사하면 모든 글자를 다 갖고 있기 때문에 이런 구분이 필요하다.

extension Trie {
  func contains(_ word: String) -> Bool {
    var currentNode = root
    var currentIndex = word.startIndex
    while currentIndex < word.endIndex {
      let c = word[currentIndex]
      guard let x = currentNode.children[c] 
      else { return false }
      currentNode = x
      currentIndex = word.index(after:currentIndex)
   }
   return currentNode.isTerminating
}

순회하기

이번에는 트리 내에 삽입된 단어들을 알파벳 순으로 순회한다. 특정 레벨에서 하위 노드의 키값을 알파벳순으로 정렬한 다음, 빠른 자식부터 방문하는 식으로 순회하면 알파벳순으로 순회할 수 있다. 다음은 완성된 단어를 발견했을 때의 처리 핸들러를 인자로 받아서 알파벳순으로 탐색한 결과에서 단어가 완성될 때 해당 핸들러를 호출하는 식으로 구현한 순회 메소드이다.

extension Trie {
  func traverse(handler: @escaping (String) -> Void)) {
    func visit(_ node: TrieNode, with characters:[Character]) {
      if node.isTerminating {
         handler(String(characters))
      }
      for c in node.children.keys.sorted() {
         let nextNode = node.children[c]!
         visit(nextNode, with: characters + [c])
      }
    }
    visit(root, with:[])
  }
}