[Data Structure] Trie

오늘은 Trie(트라이) 구조에 대해서 알아보고자 합니다.

 

Trie는 기본적으로 Tree(트리) 형태의 자료 구조입니다.

접두사 트리(Prefix Tree)라고도 하며, 주로 문자열 데이터를 다룰 때 사용됩니다.

출처: https://en.wikipedia.org/wiki/Trie

특징

- Node의 구조에 따라 다르지만, 특이한 점은 자식 Node가 이전 노드(부모 Node)의 값을 일부분을 공유

- 각 Node는 하나의 문자(or 접두어)를 저장합니다.

- 문자열의 끝을 나타내기 위한 Leaf Node 플래그를 사용합니다.

- Trie의 Root Node는 공백 문자 또는 빈 값을 저장합니다.

- 검색 및 삽입 Time Complexity는 문자열의 길이에 비례하여 O(M)입니다.(M = 문자열의 길이)

구현

1. TrieNode 구조체 정의

class TrieNode {
    var children: [Character: TrieNode] = [:]
    
    // Node의 끝 여부
    var isLeafNode: Bool = false
    
    // 해당 노드까지의 단어
    var prefix: String
    
    init(prefix: String) {
        self.prefix = prefix
    }
}

class Trie {
    // Root Node는 빈 값
    private let root = TrieNode()
}

 

2. 삽입(Insert) 구현

// in Trie Class
func insert(word: String) {
    var node = root
    var prefix = ""
    for char in word {
        if node.children[char] == nil {
            node.children[char] = TrieNode()
        }
        node = node.children[char]!
    }
    node.isLeafNode = true
}

2. 탐색

// in Trie Class
func printAllNode(node: TrieNode? = nil) {
    if let node = node {
        print(node.prefix)
    }
    for item in (node ?? root).children {
        printAllNode(node: item.value)
    }
}

3. 검색(Word & Prefix Search) 구현

// in Trie Class
func wordSearch(word: String, isPrefix: Bool) -> Bool {
    var node = root
    for char in word {
        guard let nextNode = node.children[char] else {
            return false
        }
        node = nextNode
    }
    return isPrefix ? true : node.isLeafNode
}

테스트

상단에 예제 이미지의 단어를 그대로 사용해봅시다.

let trie = Trie()
trie.insert(word: "A")
trie.insert(word: "to")
trie.insert(word: "tea")
trie.insert(word: "ted")
trie.insert(word: "ten")
trie.insert(word: "inn")
trie.printAllNode()
/*
i
in
inn
t
to
te
ten
tea
ted
A
*/

trie.search(word: "te", isPrefix: true)   // true
trie.search(word: "te", isPrefix: false)  // false
trie.search(word: "ted", isPrefix: false) // true
trie.search(word: "inn")                  // true

사용 예제

그렇다면 실제로 어떤 경우에 사용될까요?

LeetCode에서 Trie 관련 알고리즘 문제를 가져왔습니다.

https://leetcode.com/problems/longest-common-prefix

 

정확히 Trie에 대한 내용만 포함해서 Easy 난이도의 문제이며, 내용은 다음과 같습니다.

Write a function to find the longest common prefix string amongst an array of strings.
If there is no common prefix, return an empty string "".

Example 1:
  Input: strs = ["flower","flow","flight"]
  Output: "fl"
Example 2:
  Input: strs = ["dog","racecar","car"]
  Output: ""
  Explanation: There is no common prefix among the input strings.
 
Constraints:
  1 <= strs.length <= 200
  0 <= strs[i].length <= 200
  strs[i] consists of only lowercase English letters.

 

다양한 풀이가 있겠지만, 위에서 작성한 코드를 활용해보겠습니다.

class Solution {
    // isLeafNode 사용X
    class TrieNode {
        var children: [Character: TrieNode] = [:]
        var prefix: String
        
        init(prefix: String = "") {
            self.prefix = prefix
        }
    }

    func longestCommonPrefix(_ strs: [String]) -> String {
        let trie = TrieNode()
        var result = strs.first!
        for i in (0..<strs.count) {
            var node = trie
            var prefix = ""
            // 빈 문자열일 경우 바로 빈 문자열 반환
            if strs[i] == "" { return "" }
            for char in strs[i] {
                if node.children[char] == nil {
                    // 첫번째 단어는 정상적으로 insert 로직 적용
                    // 두번째 단어부터 'Node가 없음'은 다른 해당 문자까지만 Prefix가 일치함을 의미
                    if i == 0 { node.children[char] = TrieNode(prefix: prefix) }
                    else { break } 
                }
                prefix += "\(char)"
                node = node.children[char]!
            }
            // prefix와 result 중 더 짧은 값을 저장
            result = prefix.count < result.count ? prefix : result
        }
        return result
    }
}

Conclusion

오늘은 Trie에 대해 알아보고, 알고리즘 문제에 활용까지 해보았는데요.

사실 Swift에는 hasPrefix라는 함수가 이미 존재한답니다!

https://developer.apple.com/documentation/swift/string/hasprefix(_:)

hasPrefix(_:) | Apple Developer Documentation

그렇지만 Trie의 개념을 알아보고자하는 목표가 있었기 때문에!

Trie구조를 활용한 풀이와 hasPrefix를 활용한 풀이 모두 이해해보시면 좋겠습니다.