Swift 数据结构 - Tree

486 阅读2分钟

跟链表一样 Tree 也是由节点(Node)组成,每一个子节点只有一个父节点。父节点可以有多个子节点。

根节点(Root)

最顶部的节点称为根节点,根节点没有父节点。

root

叶子节点(Leaf)

没有子节点

leaf
和上一篇一样,添加一个工具方法.

func example(of desc: String,block: (()->Void)){
  print("---Example of \(desc)---")
  block()
}

ok,有了这些基础概念,我们就可以写一个TreeNode.

public class TreeNode<T> {
  public var value: T
  public var children: [TreeNode] = []
  public init(_ value: T) {
    self.value = value
  }
  public func add(_ value: TreeNode){
    children.append(value)
  }
}

创建一个简单的饮品结构。

example(of: "creating a tree") {
  let beverages = TreeNode("Beverages")

  let hot = TreeNode("Hot")
  let cold = TreeNode("Cold")

  beverages.add(hot)
  beverages.add(cold)
}

这种分层结构天然就适合使用tree

beverage

遍历一个线性的结构非常简单,但是遍历一个tree就不容易了。而且存在不同的方案。

深度优先 (depth-first)

深度优先就是从root沿一个分支探到最底部,然后再回来重新遍历。

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

就是不断递归下一个,直到没有下一个,再从头循环。

写一个完整的创建饮品tree的方法

func makeBeverageTree() -> TreeNode<String> {
  let tree = TreeNode("Beverages")

  let hot = TreeNode("hot")
  let cold = TreeNode("cold")

  let tea = TreeNode("tea")
  let coffee = TreeNode("coffee")
  let chocolate = TreeNode("cocoa")

  let blackTea = TreeNode("black")
  let greenTea = TreeNode("green")
  let chaiTea = TreeNode("chai")

  let soda = TreeNode("soda")
  let milk = TreeNode("milk")

  let gingerAle = TreeNode("ginger ale")
  let bitterLemon = TreeNode("bitter lemon")

  tree.add(hot)
  tree.add(cold)

  hot.add(tea)
  hot.add(coffee)
  hot.add(chocolate)

  cold.add(soda)
  cold.add(milk)
  
  tea.add(blackTea)
  tea.add(greenTea)
  tea.add(chaiTea)
  
  soda.add(gingerAle)
  soda.add(bitterLemon)
  return tree
}

这个树形结构大概长这样。

tree

利用上面的depth-first策略进行遍历。

example(of: "depth-first traversal") {
  let tree = makeBeverageTree()
  tree.forEachDepthFirst{ print($0.value) }
}

// result:
Beverages
hot
tea
black
green
chai
coffee
cocoa
cold
soda
ginger ale
bitter lemon
milk

符合预期。

层级顺序(level-order)

顾名思义也就是一层一层遍历

extension TreeNode {
  public 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) })
    }
  }
}

这里用到了队列的特性,FIFO(先进先出)。按层级入队。

一个简单的队列结构Queue

public struct Queue<T> {
  public mutating func enqueue(_ element: T) {
    array.append(element)
  }
  public mutating func dequeue() -> T? {
    return isEmpty ? nil:array.removeFirst()
  }
  public var isEmpty: Bool {
    return array.isEmpty
  }
  public var peek: T? {
    return array.first
  }
  public typealias Element = T
  private var array: [T] = []
  public init() {}
}
example(of: "level-order traversal") {
  let tree = makeBeverageTree()
  tree.forEachLevelOrder { print($0.value) }
}
// result: 
---Example of level-order traversal---
Beverages
hot
cold
tea
coffee
cocoa
soda
milk
black
green
chai
ginger ale
bitter lemon

符合预期 确实是按层级顺序输出

根据value查询节点

查询就比较简单了,因为我们已经写好了遍历的方法。

extension TreeNode where T: Equatable {
  public func search(_ value: T) -> TreeNode? {
    var result: TreeNode?
    forEachLevelOrder { (node) in
      if node.value == value {
        result = node
      }
    }
    return result
  }
}

sample:

example(of: "searching for a node") {
  let tree = makeBeverageTree()
  if let searchResult1 = tree.search("ginger ale") {
    print("Found node: \(searchResult1.value)")
  }
  
  if let searchResult2 = tree.search("WKD Blue") {
    print(searchResult2.value)
  } else {
    print("Couldn't find WKD Blue")
  }
}

// result: 
---Example of searching for a node---
Found node: ginger ale
Couldn't find WKD Blue

按level分行输出

我们想把同一个level的放在一行输出。 同样用到队列的特性,对每行使用一个队列,然后行数进行递归。

extension TreeNode {
  
  private func traversalSameLine(_ queue: Queue<TreeNode>){
    var queue = queue
    var nextQueue = Queue<TreeNode>()
    var values = ""
    while let node = queue.dequeue() {
      values += "\(node.value) "
      node.children.forEach({ nextQueue.enqueue($0) })
    }
    print("\(values)\n")
    if !nextQueue.isEmpty {
      traversalSameLine(nextQueue)
    }
  }
  public func printForLine(){
    print("\(value)\n")
    var queue = Queue<TreeNode>()
    children.forEach({ queue.enqueue($0) })
    traversalSameLine(queue)
  }
}
example(of: "traversal in line") {
  let tree = makeBeverageTree()
  tree.printForLine()
}
// result:
---Example of traversal in line---
Beverages

hot cold 

tea coffee cocoa soda milk 

black green chai ginger ale bitter lemon 

结果符合预期。


内容来自《Data Structure & Algorithms in Swift》简记