数据结构学习-Trie

介绍

Trie （发音 try）又可以叫成单词树，专门用于存储可以表示为集合的数据，比如英语单词。数据结构如下：

每个字母都是一个节点，其中有一个小点标示的是结束节点，比如图中的 g， r 等字母。

为什么使用它？假如要查询一个搜索词的所有匹配项。如果使用数组存储的话，那么搜索的时间复杂度是 O(k * n)， K: 最大单词长度，n: 数组长度。而如果是Trie数据结构的话，那么时间复杂度是 O(k * m), m: 匹配的结果数目。

比如搜索有 car 字符串开头的单词，那么满足条件的有： car，care， card， cargo。

实现

节点

（下面以英文单词的插入与查找为例子描述）

和所有的树结构一样，得有一个节点存储数据信息。下面是 TrieNode 的实现：

class TrieNode<Key: Hashable> {
    
    // 当前存储的值，比如存 a 字母
    var key: Key?
    
    // 父节点的弱引用
    weak var parent: TrieNode?
    
    // 包含的子节点字典
    var children: [Key: TrieNode] = [:]
    
    // 是否是结束节点
    var isTerminating = false
    
    init(key: Key?, parent: TrieNode?) {
        self.key = key
        self.parent = parent
    }
}

Trie

Trie 树的基本实现：

class Trie<CollectionType: Collection & Hashable> where CollectionType.Element: Hashable {
    
    typealias Node = TrieNode<CollectionType.Element>
    
    // 根节点
    private let root = Node(key: nil, parent: nil)
    
    // 包含的所有单词
    public private(set) var collections: Set<CollectionType> = []
}

根节点是数据的入口，里面不存数据。

插入

extension Trie {
    
    func insert(_ collection: CollectionType) {
        var current = root
        
        for element in collection {
            // 如果子节点子没有这个字母，则创建
            if current.children[element] == nil {
                current.children[element] = Node(key: element, parent: current)
            }
            current = current.children[element]!
        }
        
        // 如果不是重复插入的单词 最后标记为结束节点
        if current.isTerminating {
            return
        } else {
            // 标记
            current.isTerminating = true
            
            // 添加单词（collections是本地的一个变量，不影响 Trie 树的结构）
            collections.insert(collection)
        }
    }
}

插入操作时间复杂度 O(k), k: 插入的数据长度。

对于上面的实现代码，以 car 单词的插入为例。由于目前Trie树里面为空，所以插入后的 Trie 树是这样的：

插入单词 card 后：

插入单词 cargo 后：

插入单词 dog 后：

包含判断

extension Trie {

    func contains(_ collection: CollectionType) -> Bool {
        var current = root
        for element in collection {
            guard let child = current.children[element] else {
                return false
            }
            current = child
        }
        return current.isTerminating
    }
}

包含查询的时间复杂度为O(k), k: 查找的数据长度。

代码的实现和插入类似，都是查找子节点，看匹配情况。

删除

删除操作需要注意的要点是该删除数据的结尾处是在 Trie 树的中间节点还是叶子节点。代码实现：

extension Trie {

    func remove(_ collection: CollectionType) {
        var current = root
        for element in collection {
            guard let child = current.children[element] else {
                // 不匹配，直接返回
                return
            }
            current = child
        }
        guard current.isTerminating else {
            // 不是完整单词，不修改树的结构， 直接返回
            return
        }
        
        // 标记为非单词结尾
        current.isTerminating = false
        
        // 移除单词（collections是本地的一个变量，不影响 Trie 树的结构）
        collections.remove(collection)

        // 如果必要，递归移除
        while let parent = current.parent, current.children.isEmpty && !current.isTerminating {
            parent.children[current.key!] = nil
            current = parent
        }
    }
}

移除操作的时间复杂度为O(k), k: 数据长度。

测试：

    let trie = Trie<String>()
    trie.insert("car")
    trie.insert("card")
    
    print("\n*** 移除 car 前 ***")
    print("所有单词：\(trie.collections)")

    print("\n*** 移除 car 后 ***")
    trie.remove("car")
    print("所有单词：\(trie.collections)")

结果：

*** 移除 car 前 ***
所有单词：["car", "card"]

*** 移除 car 后 ***
所有单词：["card"]

前缀匹配查询

extension Trie where CollectionType: RangeReplaceableCollection {
    
    func collections(startingWith prefix: CollectionType) -> [CollectionType] {
        var current = root
        
        // 判断查询有效
        for element in prefix {
            guard let child = current.children[element] else {
                return []
            }
            current = child
        }
        
        return collections(startingWith: prefix, after: current)
    }
    
    func collections(startingWith prefix: CollectionType, after node: Node) -> [CollectionType] {
        var results: [CollectionType] = []
        
        // 当前是结束位置
        if node.isTerminating {
            results.append(prefix)
        }
        
        // 遍历所有子节点
        for child in node.children.values {
            var prefix = prefix
            prefix.append(child.key!)
            results.append(contentsOf: collections(startingWith: prefix, after: child))
        }
        
        return results
    }
}

测试：

    let trie = Trie<String>()
    trie.insert("car")
    trie.insert("card")
    trie.insert("care")
    trie.insert("cared")
    trie.insert("cars")
    trie.insert("carbs")
    trie.insert("carapace")
    trie.insert("cargo")
    trie.insert("dog")
    trie.insert("apple")

    print("所有单词：\(trie.collections)")

    print("\n有 car 前缀单词")
    let prefixedWithCar = trie.collections(startingWith: "car")
    print(prefixedWithCar)
    
    print("\n有 care 前缀单词")
    let prefixedWithCare = trie.collections(startingWith: "care")
    print(prefixedWithCare)

结果：

所有单词：["carapace", "dog", "cars", "apple", "care", "carbs", "cargo", "cared", "card", "car"]

有 car 前缀单词
["car", "cars", "carapace", "cargo", "card", "carbs", "care", "cared"]

有 care 前缀单词
["care", "cared"]