数据结构学习-树(Trees)

之前学习的栈，链表，队列都是线性数据结构。而树形数据结构则在代表层级关系（比如目录），管理已排序数据，加快数据查找速度等有着大量的应用。

一些术语

和链表一样，树也是使用节点（Node）代表里面的元素的。

父节点，子节点

下图是一个常见的树，它的访问位置是从上到下的（一棵实际生活中倒序的树）

在树中，每个节点都有一个上面的节点（除了最顶部那个节点），上面的那个节点称为父节点。子节点只有一个父节点。

根节点

如图红色标示所示，根节点是树的开始位置，它没有父节点。

叶子节点

如图红色标示所示，叶子节点在树中是没有子节点的。

实现

由于和链表一样，它们的元素都是节点，下面添加树节点实现：

class TreeNode<T> {
    var value: T
    var children: [TreeNode] = [] // 保存所有的子节点（可为空）
    
    init(value: T) {
        self.value = value
    }
    
    func add(_ child: TreeNode<T>) {
        children.append(child)
    }
}

添加测试代码：

    let animal = TreeNode(value: "animal")
    animal.add(TreeNode(value: "cat"))
    animal.add(TreeNode(value: "dog"))

它的树结构是这样的：

遍历

深度优先

深度优先是先访问根节点，再访问它的子节点，再访问子节点的子节点，直到叶子节点。

下面是算法实现：

extension TreeNode {
    
    func forEachDepthFirst(visit: (TreeNode) -> Void) {
        visit(self)
        children.forEach({
            $0.forEachDepthFirst(visit: visit)
        })
    }
}

创建一棵树的辅助代码：

func createAnimals() -> TreeNode<String> {
    let animal = TreeNode(value: "animal")
    
    let cat = TreeNode(value: "cat")
    let dog = TreeNode(value: "dog")
    
    let black = TreeNode(value: "black")
    let white = TreeNode(value: "white")
    
    let big = TreeNode(value: "big")
    let middle = TreeNode(value: "middle")
    let small = TreeNode(value: "small")
    
    let brutal = TreeNode(value: "brutal")
    let cute = TreeNode(value: "cute")
    
    animal.add(cat)
    animal.add(dog)
    
    cat.add(black)
    cat.add(white)
    
    dog.add(big)
    dog.add(middle)
    dog.add(small)
    
    big.add(brutal)
    big.add(cute)
    
    return animal

它的树结构是这样的：

代码测试：

func depthFirstVisit() {
    let animal = createAnimals()
    animal.forEachDepthFirst(visit: { print($0.value) })
}

输出：

animal
cat
black
white
dog
big
brutal
cute
middle
small

层级优先

层级优先是先访问所有同一层级的节点。

下面是算法实现：

extension TreeNode {

    func forEachLevelOrder(visit: (TreeNode) -> Void) {
        visit(self)
        var queue = Queue<TreeNode>() // 使用队列保存当前层级的所有节点
        children.forEach({ queue.enqueue($0) })
        while let node = queue.dequeue() {
            visit(node)
            node.children.forEach({ queue.enqueue($0) })
        }
    }
}

算法实现的关键点是使用队列保存当前层级的所有节点，队列实现这里有介绍，下面直接把实现代码贴过来：

struct Queue<T> {
    
    private var leftStack: [T] = []
    private var rightStack: [T] = []
    
    init() {}
    
    var isEmpty: Bool {
        leftStack.isEmpty && rightStack.isEmpty
    }
    
    var peek: T? {
        !leftStack.isEmpty ? leftStack.last : rightStack.first
    }
    
    var count: Int {
        leftStack.count + rightStack.count
    }
    
    @discardableResult mutating func enqueue(_ element: T) -> Bool {
        rightStack.append(element)
        return true
    }
    
    mutating func dequeue() -> T? {
        if leftStack.isEmpty {
            leftStack = rightStack.reversed()
            rightStack.removeAll()
        }
        return leftStack.popLast()
    }
}

代码测试：

func levelVisit() {
    let animal = createAnimals()
    animal.forEachLevelOrder(visit: { print($0.value) })
}

输出：

animal
cat
dog
black
white
big
middle
small
brutal
cute

查找

由于树已经实现了深度优先和层级优先遍历算法，所以查找树中的某个节点其实也可以基于不同的遍历算法实现。

实现：

extension TreeNode where T: Equatable {

    func search(_ value: T) -> TreeNode? {
        var result: TreeNode?
        
        // 使用层级优先遍历算法查找，目前是找到最后一个相同的节点
        forEachLevelOrder { node in
            if node.value == value {
                result = node
            }
        }
        return result
    }
}

练习

按层级输出, 比如上面的animal树，应该输出这样：

animal

cat dog

black white big middle small

brutal cute

第一种代码实现：

func printInLevelOrder<T>(items: [TreeNode<T>]) {
    print(items.map({ "\($0.value)" }).joined(separator: " "))
    
    var results: [TreeNode<T>] = []
    for i in items {
        i.children.forEach({ results.append($0) })
    }
    
    if !results.isEmpty {
        printInLevelOrder(items: results)
    }
}

测试代码：

    let animal = createAnimals()
    printInLevelOrder(items: [animal])

可以看到输出是OK的。

第二种代码实现：

func printEachLevel<T>(for tree: TreeNode<T>) {
    var queue = Queue<TreeNode<T>>()
    var nodesLeftInCurrentLevel = 0
    queue.enqueue(tree)
    
    while !queue.isEmpty {
        nodesLeftInCurrentLevel = queue.count
        
        while nodesLeftInCurrentLevel > 0 {
            guard let node = queue.dequeue() else { break }
            print("\(node.value) ", terminator: "")
            node.children.forEach { queue.enqueue($0) }
            nodesLeftInCurrentLevel -= 1
        }
        
        print()
    }
}

第二种实现的关键是使用队列保存当前层级的数目。