Swift实现二叉搜索树与字典树

二叉搜索树

二叉搜索树的性质如下：

• 若任意节点的左子树不空，则左子树所有节点的值小于根节点的值
• 若任意节点的右子树不空，则左子树所有节点的值大于根节点的值
• 任意节点的左右子树也为二叉搜索树
• 没有键值相等的节点

其基本操作有如下四个

• empty——返回一个空的二叉搜索树
• isEmpty——检查二叉搜索树是否为空
• contains——检查二叉搜索树是否包含某元素
• insert——向二叉搜索树插入元素

具体实现过程

• 用一个递归枚举定义二叉搜索树的数据结构

indirect enum BinarySearchTree<E:Comparable>{ // indirect 表明是递归枚举
    case Leaf
    case Node(BinarySearchTree<E>, E, BinarySearchTree<E>)
}

通过关联值指定一个节点的左右子树及其自身。

• 初始化方法，可以初始化一个空树或者一个根节点有值的树

init(){
    self = .Leaf
}

init(value:E) {
    self = .Node(.Leaf, value, .Leaf)
}

• 打印所有元素的方法

var elements:[E] {
    switch self {
    case .Leaf:
        return []
    case let .Node(left, x, right):
        return left.elements + [x] + right.elements
    }
}

这个方法可以按照中序排列方式打印出所有元素

• 判断是否为空以及是否为二叉搜索树的方法

var isEmpty:Bool {
    if case .Leaf = self {
        return true
    }
    return false
}

var isBST:Bool {
    switch self {
    case .Leaf:
        return true
    case let .Node(left, x, right):
        return left.elements.reduce(true, { (result, y) -> Bool in
            result && (y < x)
        }) && right.elements.reduce(true, { (result, y) -> Bool in
            result && (y > x)
        }) && left.isBST && right.isBST
    }
}

其中 isBST 方法会根据二叉搜索树的前三个特性来检测一个树。

• 检测是否包含某元素的方法

func contains(x:E) -> Bool {
    switch self {
    case .Leaf:
        return false
    case let .Node(_, y, _) where x == y:
        return true
    case let .Node(left, y, _) where x < y:
        return left.contains(x: x)
    case let .Node(_, y, right) where x > y:
        return right.contains(x: x)
    default:
        fatalError()
    }
}

如果树是空的，则 x 不在树中，返回 false。

如果树不为空，且储存在根节点的值与 x 相等，返回 true。

如果树不为空，且储存在根节点的值大于 x，那么如果 x 在树中的话，它一定是在左子树中，所以，我们在左子树中递归搜索 x。

类似地，如果根节点的值小于 x，我们就在右子树中继续搜索。

• 插入一个元素

mutating func insert(x:E) {
    switch self {
    case .Leaf:
        self = BinarySearchTree.init(value: x)
    case .Node(var left, let y, var right):
        if x < y {left.insert(x: x)}
        if x > y {right.insert(x: x)}
        self = .Node(left, y, right)
    }
}

由于需要对二叉搜索树进行修改，所以该方法用到了 mutating 来修饰。方法递归遍历二叉搜索树，直到找到一个空叶子节点就加入。

使用如下

var binaryTree = BinarySearchTree.init(value: 5)
binaryTree.insert(x: 1)
binaryTree.insert(x: 2)
binaryTree.insert(x: 3)
binaryTree.insert(x: 4)
binaryTree.insert(x: 5)
binaryTree.insert(x: 6)
binaryTree.insert(x: 7)
binaryTree.insert(x: 8)
binaryTree.insert(x: 9)
binaryTree.insert(x: 10)
print(binaryTree.elements)
// [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

字典树

字典树 Trie 是一种有序树，用于保存关联数组，通常是字符串，字典树的键不是直接保存在节点里，而是由一串节点的值组合而成。

具体实现过程

• 用结构体定义一个字典树的数据结构

struct Trie<Element:Hashable>{
    var children:[Element:Trie]
    var isElement: Bool
}

其中 children 就是一个字典，值是一个子字典树，键是子树的前缀值，而 isElement 定义了当前节点为止的组合是否是一个元素

例如上图中，“at” 节点的 isElement 就是 true，而 ag 由于不是一个有意义的单词所以 isElement 为 false。

• 初始化方法

首先是两个基础的初始化方法

init() {
    isElement = false
    children = [:]
}

init(isElement:Bool, children:[Element:Trie]) {
    self.isElement = isElement;
    self.children = children;
}

但是正常使用过程中，我们更需要一个能接受一个数组作为入参的初始化方法，然后形成一个单元素的字典树。

此时我们首先需要实现一个数组的递归方式，即从最原始的数组逐渐去除一定数量的元素形成子数组，我们将其定义为 Array 的一个拓展方法

extension Array {
    var decompose: (Element, [Element])? {
        return isEmpty ? nil : (self[startIndex], Array(self.dropFirst()))
    }
}

可以看到，如果数组为空则返回 nil，否则返回一个元组，元组第一个元素是当前数组的头元素，第二个元素是去掉第一个元素之后由该数组其余元素组成的新数组。我们可以通过重复调用一个数组的 decompose 方法递归地遍历一个数组。

接下来就可以定义一个接收数组参数的初始化方法了

init(key:[Element]) {
    if let (head, tail) = key.decompose {
        let children = [head: Trie.init(key: tail)]
        self = Trie(isElement: false, children: children)
    } else {
        self = Trie(isElement: true, children: [:])
    }
}

可以看到它遍历传入的数组，生成以数组中元素为节点的单元素字典树。

• 插入操作

func insert(key:[Element]) -> Trie<Element> {
    guard let (head, tail) = key.decompose else {
        return Trie.init(isElement: true, children: self.children)
    }
    var newChildren = self.children
    if let nextTrie = self.children[head] {
        newChildren[head] = nextTrie.insert(key: tail)
    } else {
        newChildren[head] = Trie.init(key: tail)
    }
    return Trie.init(isElement: isElement, children: newChildren)
}

遍历传入的数组并检查每一个节点的 children 键组。

如果键组为空，我们将 isElement 设置为 true，然后不再修改剩余的字典树，表明已经插入完毕了。

如果键组不为空，且键组的 head 已经存在于当前节点的 children 字典中，我们只需要递归地调用该函数，将键组的 tail 插入到对应的子字典树中。

如果键组不为空，且第一个键 head 并不是该字典树中 children 字典的某条记录，就创建一棵新的字典树来储存键组中剩下的键。然后，以 head 键和 tail 数组生成对应新的字典树，储存在当前节点中，完成插入操作。

然后我们就可以定义一个方法来将一个关联数组组合成字典树

func buildStringTrie(words:[String]) -> Trie<Character> {
    let emptyTrie = Trie<Character>.init()
    return words.reduce(emptyTrie, { (result, word) -> Trie<Character> in
        result.insert(key: Array.init(word))
    })
}

• 打印方法，打印出字典树中包含的所有关联元素

var elements: [[Element]] {
    var result: [[Element]] = isElement ? [[]] : []
    for (key, value) in children {
        result += value.elements.map { [key] + $0 }
    }
    return result
}

这个方法很精巧，同时也很难懂，首先需要注意下面一些操作的结果

[] + [["r"]] = [["r"]]
[[]] + [["r"]] = [[], ["r"]]

对于一个嵌套的二维数组，直接加一个空的一维数组，这个空的一维数组就被加入为二维数组的一个元素，而两个嵌套的二维数组相加，则会合并为一个二维数组。

现在来看这段代码，它的思路其实是：

• 对于 isElement 为 true 的节点，返回类似 [[]] 或 [[], [“X”]] 的数组，里面一定包含一个空数组，开始将此节点回溯的所有父节点组合为元素
• 对于 isElement 为 false 的节点，由于它一定是某个 isElement 为 true 的节点的父节点，所以返回类似 [[“X”], [“Y”, “Z”]] 的数组，不会生成多余的数组来组合新的元素
• 依次遍历，最终生成包含所有元素的数组

让我们用一个简单的例子来测试，首先用 [“cart”, “car”] 数组来生成一个字典树，然后在 elements 过程中进行打印

{
        let contents = ["cart", "car"]
        let trie = buildStringTrie(words:contents)
        print(trie.elements)
}

{
    var elements: [[Element]] {
        var result: [[Element]] = isElement ? [[]] : []
        for (key, value) in children {
            result += value.elements.map { [key] + $0 }
        }
        print(result)
        return result
    }
}

最终结果如下

[[]]
[[], ["t"]]
[["r"], ["r", "t"]]
[["a", "r"], ["a", "r", "t"]]
[["c", "a", "r"], ["c", "a", "r", "t"]]
[["c", "a", "r"], ["c", "a", "r", "t"]]

• 检索某个字符串是否在字典树的方法和检索某个前缀对应的字典树的方法

func lookup(key: [Element]) -> Bool {
    guard let (head, tail) = key.decompose else { return isElement }
    guard let subtrie = children[head] else { return false }
    return subtrie.lookup(key: tail)
}

func withPrefix(prefix:[Element]) -> Trie<Element>?{
    guard let (head, tail) = prefix.decompose else {return self}
    guard let remainder = children[head] else {return nil}
    return remainder.withPrefix(prefix: tail)
}

思路大致一样，遍历传入的数组，当遍历结束时返回结果。