二叉树

二叉树 Swift 实现

import Foundation

enum TraverseType {
    case preOrder
    case inOrder
    case postOrder
}

class BinaryTree<T> {
    var data: T
    var leftChild: BinaryTree<T>?
    var rightChild: BinaryTree<T>?
    
    init(data: T, leftChild: BinaryTree<T>? = nil, rightChild: BinaryTree<T>? = nil) {
        self.data = data
        self.leftChild = leftChild
        self.rightChild = rightChild
    }
}

extension BinaryTree {
    var count: Int {
        get {
            var count = 0
            traverse { _ in count += 1 }
            return count
        }
    }
    
    var isEmpty: Bool {
        get {
            return (leftChild == nil && rightChild == nil)
        }
    }
    
    var depth: Int {
        get {
            // To-do
            return 0
        }
    }
}

extension BinaryTree {
    typealias TraverseCallback = (T) -> Void
    
    func traverse(type: TraverseType = .preOrder, callback: TraverseCallback) {
        BinaryTree.traverse(tree: self, type: type, callback: callback)
    }
    
    static func traverse(tree: BinaryTree?,
                         type: TraverseType = .preOrder, callback: TraverseCallback) {
        guard let tree = tree else { return }
        
        if type == .preOrder { callback(tree.data) }
        traverse(tree: tree.leftChild, callback: callback)
        if type == .inOrder { callback(tree.data) }
        traverse(tree: tree.rightChild, callback: callback)
        if type == .postOrder { callback(tree.data) }
    }
}




let node1 = BinaryTree(data: "A")
let node2 = BinaryTree(data: "B")
let node3 = BinaryTree(data: "C")
let node4 = BinaryTree(data: "D")
let node5 = BinaryTree(data: "E")
let node6 = BinaryTree(data: "F")

let root = BinaryTree(data: "ROOT", leftChild: node1, rightChild: node2)

node1.leftChild = node3
node1.rightChild = node4

node2.leftChild = node5
node2.rightChild = node6


root.traverse(type: .preOrder) { print($0) }

root.traverse(type: .inOrder) { print($0) }

root.traverse(type: .postOrder) { print($0) }