How to draw a binary tree in Swift 4?

Question

Based on this Ray Wenderlich article I'm able to create a binary tree data structure like that:

  enum BinaryTree<T: Comparable> {

  case empty
  indirect case node(BinaryTree<T>, T, BinaryTree<T>)

  var count: Int {
    switch self {
    case let .node(left, _, right):
      return left.count + 1 + right.count
    case .empty:
      return 0
    }
  }

  // 1.
  mutating func naiveInsert(newValue: T) {
    // 2.
    guard case .node(var left, let value, var right) = self else {
      // 3. 
      self = .node(.empty, newValue, .empty)
      return 
    }

    // 4. TODO: Implement naive algorithm!
    if newValue < value {
      left.naiveInsert(newValue: newValue)
    } else {
      right.naiveInsert(newValue: newValue)
    }

  }

  private func newTreeWithInsertedValue(newValue: T) -> BinaryTree {
    switch self {
    // 1
    case .empty:
      return .node(.empty, newValue, .empty)
    // 2 
    case let .node(left, value, right):
      if newValue < value {
        return .node(left.newTreeWithInsertedValue(newValue: newValue), value, right)
      } else {
        return .node(left, value, right.newTreeWithInsertedValue(newValue: newValue))
      }
    }
  }

  mutating func insert(newValue: T) {
    self = newTreeWithInsertedValue(newValue: newValue)
  }

    func traverseInOrder(process: (T) -> ()) {
    switch self {
    // 1
    case .empty:
      return 
    // 2
    case let .node(left, value, right):
      left.traverseInOrder(process: process)
      process(value)
      right.traverseInOrder(process: process)
    }
  }

    func traversePreOrder( process: (T) -> ()) {
    switch self {
    case .empty:
      return
    case let .node(left, value, right):
      process(value)
      left.traversePreOrder(process: process)
      right.traversePreOrder(process: process)
    }
  }

    func traversePostOrder( process: (T) -> ()) {
    switch self {
    case .empty:
      return
    case let .node(left, value, right):
      left.traversePostOrder(process: process)
      right.traversePostOrder(process: process)
      process(value) 
    }
  }

  func search(searchValue: T) -> BinaryTree? {
    switch self {
    case .empty:
      return nil
    case let .node(left, value, right):
      // 1
      if searchValue == value {
        return self
      }

      // 2
      if searchValue < value {
        return left.search(searchValue: searchValue)
      } else {
        return right.search(searchValue: searchValue)
      }
    }
  }

}

extension BinaryTree: CustomStringConvertible {
  var description: String {
    switch self {
    case let .node(left, value, right):
      return "value: \(value), left = [" + left.description + "], right = [" + right.description + "]"
    case .empty:
      return ""
    }
  }
}

// leaf nodes
let node5 = BinaryTree.node(.empty, "5", .empty)
let nodeA = BinaryTree.node(.empty, "a", .empty)
let node10 = BinaryTree.node(.empty, "10", .empty)
let node4 = BinaryTree.node(.empty, "4", .empty)
let node3 = BinaryTree.node(.empty, "3", .empty)
let nodeB = BinaryTree.node(.empty, "b", .empty)

// intermediate nodes on the left
let Aminus10 = BinaryTree.node(nodeA, "-", node10)
let timesLeft = BinaryTree.node(node5, "*", Aminus10)

// intermediate nodes on the right
let minus4 = BinaryTree.node(.empty, "-", node4)
let divide3andB = BinaryTree.node(node3, "/", nodeB)
let timesRight = BinaryTree.node(minus4, "*", divide3andB)

// root node
var tree: BinaryTree<Int> = .empty
tree.insert(newValue: 7)
tree.insert(newValue: 10)
tree.insert(newValue: 2)
tree.insert(newValue: 1)
tree.insert(newValue: 5)
tree.insert(newValue: 9)
tree.insert(newValue: 3)

tree.traverseInOrder { print($0) }
tree.search(searchValue: 5)

I found a lot of examples here on stack to visualize such a tree in Android Graphical binary tree in Android or PHP draw binary tree with php but nothing in Swift. I thought about the Core Graphics library but where to start? Can anyone give me an example?

Answer 1

In terms of the basics of how to draw a line, you:

1. Create UIBezierPath;
2. Move to the starting point with move(to:);
3. Adding line to the end point with addLine(to:);

You can then render that path in your UI by either:

• Create CAShapeLayer, specifying its strokeWidth, strokeColor, and fillColor; set its path, and then add that shape layer as a sublayer of the view’s layer; or
• Create UIView subclass and in its draw(_:) method you can call setStroke of the desired UIColor, set the lineWidth of the UIBezierPath, and then stroke() the UIBezierPath.

Generally, I’d use the CAShapeLayer approach, where I basically configure the shape layer, but let the OS render that shape layer for me.

---

That having been said, I’d probably take this a step further, and wrap the line drawing in its own UIView subclass. The thought process is that not only are high-level views generally composed of UIView objects, but it also opens the door for all sorts of advanced UX (e.g. you might want to detect a tap on a node and do something).

Anyway, I’d wrap that “connector line” drawing code in a UIView subclass, like so:

class ConnectorView: UIView {
    enum ConnectorType {
        case upperRightToLowerLeft
        case upperLeftToLowerRight
        case vertical
    }

    var connectorType: ConnectorType = .upperLeftToLowerRight { didSet { layoutIfNeeded() } }

    override class var layerClass: AnyClass { return CAShapeLayer.self }
    var shapeLayer: CAShapeLayer { return layer as! CAShapeLayer }

    convenience init(connectorType: ConnectorType) {
        self.init()
        self.connectorType = connectorType
    }

    override init(frame: CGRect = .zero) {
        super.init(frame: frame)
        configure()
    }

    required init?(coder aDecoder: NSCoder) {
        super.init(coder: aDecoder)
        configure()
    }

    override func layoutSubviews() {
        let path = UIBezierPath()

        switch connectorType {
        case .upperLeftToLowerRight:
            path.move(to: CGPoint(x: bounds.minX, y: bounds.minY))
            path.addLine(to: CGPoint(x: bounds.maxX, y: bounds.maxY))

        case .upperRightToLowerLeft:
            path.move(to: CGPoint(x: bounds.maxX, y: bounds.minY))
            path.addLine(to: CGPoint(x: bounds.minX, y: bounds.maxY))

        case .vertical:
            path.move(to: CGPoint(x: bounds.midX, y: bounds.minY))
            path.addLine(to: CGPoint(x: bounds.midX, y: bounds.maxY))
        }

        shapeLayer.path = path.cgPath
    }

    override var description: String { return String(format: "<ConnectorView: %p; frame = %@, type = %@", self, frame.debugDescription, connectorType.string) }
}

private extension ConnectorView {
    func configure() {
        shapeLayer.lineWidth = 3
        shapeLayer.strokeColor = UIColor.black.cgColor
        shapeLayer.fillColor = UIColor.clear.cgColor
    }
}

This defines the shape layer to stroke the line from one corner to another, which it will automatically update accordingly as the frame of this view changes. By doing this, I can now control where the connector line view is rendered by updating a frame of this UIView subclass. The virtue of this approach is that I can now define the constraints for this ConnectorView such that the top/bottom/left/right anchors are tied to the centerX and centerY of the UIView for the respective nodes. And by putting the nodes in front of these connector line views, it results in the desired look and feel.

FYI, for simple rectangular nodes, you might just subclassed UILabel for the nodes, themselves:

class NodeView: UILabel {
    weak var containerView: UIView!

    override init(frame: CGRect = .zero) {
        super.init(frame: frame)
        configure()
    }

    required init?(coder aDecoder: NSCoder) {
        super.init(coder: aDecoder)
        configure()
    }
}

private extension NodeView {
    func configure() {
        backgroundColor = UIColor.white
        layer.borderColor = UIColor.black.cgColor
        layer.borderWidth = 3
        textAlignment = .center
    }
}

Now, the trick is where to place the nodes so that there’s enough space for all their child nodes. If you’re new to the iOS constraints system this is going to look super confusing (frankly, even if you are familiar with it, it’s a bit ugly), but you can do something like:

private let nodeSpacing: CGFloat = 50
private let nodeVerticalSpace: CGFloat = 50
private let nodeHorizontalSpace: CGFloat = 50
private let nodeHeight: CGFloat = 40
private let nodeWidth: CGFloat = 60

extension BinaryTree {

    func addNodes(to view: UIView) -> NodeView? {
        guard case .node(let leftNode, let value, let rightNode) = self else { return nil }

        let containerView = UIView()
        containerView.translatesAutoresizingMaskIntoConstraints = false
        view.addSubview(containerView)

        let thisNodeView = NodeView()
        thisNodeView.translatesAutoresizingMaskIntoConstraints = false
        thisNodeView.text = String(describing: value)
        thisNodeView.containerView = containerView
        containerView.addSubview(thisNodeView)

        NSLayoutConstraint.activate([
            containerView.topAnchor.constraint(equalTo: thisNodeView.topAnchor),

            thisNodeView.widthAnchor.constraint(equalToConstant: nodeWidth),
            thisNodeView.heightAnchor.constraint(equalToConstant: nodeHeight),
        ])

        switch (leftNode, rightNode) {
        case (.empty, .empty):
            NSLayoutConstraint.activate([
                containerView.bottomAnchor.constraint(equalTo: thisNodeView.bottomAnchor),
                containerView.leftAnchor.constraint(equalTo: thisNodeView.leftAnchor),
                containerView.rightAnchor.constraint(equalTo: thisNodeView.rightAnchor)
            ])

        case (let node, .empty), (.empty, let node):
            let nodeView = node.addNodes(to: containerView)!
            let connector = ConnectorView(connectorType: .vertical)
            connector.translatesAutoresizingMaskIntoConstraints = false
            containerView.insertSubview(connector, belowSubview: thisNodeView)

            NSLayoutConstraint.activate([
                thisNodeView.bottomAnchor.constraint(equalTo: nodeView.topAnchor, constant: -nodeVerticalSpace),
                thisNodeView.centerXAnchor.constraint(equalTo: nodeView.centerXAnchor),
                connector.topAnchor.constraint(equalTo: thisNodeView.centerYAnchor),
                connector.bottomAnchor.constraint(equalTo: nodeView.centerYAnchor),
                connector.leadingAnchor.constraint(equalTo: thisNodeView.leadingAnchor),
                connector.trailingAnchor.constraint(equalTo: thisNodeView.trailingAnchor),

                containerView.bottomAnchor.constraint(equalTo: nodeView.containerView.bottomAnchor),
                containerView.leftAnchor.constraint(equalTo: nodeView.containerView.leftAnchor),
                containerView.rightAnchor.constraint(equalTo: nodeView.containerView.rightAnchor)
            ])

        case (let leftNode, let rightNode):
            let leftNodeView = leftNode.addNodes(to: containerView)!
            let rightNodeView = rightNode.addNodes(to: containerView)!

            let leftConnector = ConnectorView(connectorType: .upperRightToLowerLeft)
            leftConnector.translatesAutoresizingMaskIntoConstraints = false
            containerView.insertSubview(leftConnector, belowSubview: thisNodeView)

            let rightConnector = ConnectorView(connectorType: .upperLeftToLowerRight)
            rightConnector.translatesAutoresizingMaskIntoConstraints = false
            containerView.insertSubview(rightConnector, belowSubview: thisNodeView)

            for nodeView in [leftNodeView, rightNodeView] {
                NSLayoutConstraint.activate([
                    thisNodeView.bottomAnchor.constraint(equalTo: nodeView.topAnchor, constant: -nodeVerticalSpace),
                ])
            }
            NSLayoutConstraint.activate([
                leftNodeView.containerView.rightAnchor.constraint(lessThanOrEqualTo: rightNodeView.containerView.leftAnchor, constant: -nodeHorizontalSpace),

                leftConnector.topAnchor.constraint(equalTo: thisNodeView.centerYAnchor),
                leftConnector.bottomAnchor.constraint(equalTo: leftNodeView.centerYAnchor),
                leftConnector.leadingAnchor.constraint(equalTo: leftNodeView.centerXAnchor),
                leftConnector.trailingAnchor.constraint(equalTo: thisNodeView.centerXAnchor),

                rightConnector.topAnchor.constraint(equalTo: thisNodeView.centerYAnchor),
                rightConnector.bottomAnchor.constraint(equalTo: rightNodeView.centerYAnchor),
                rightConnector.leadingAnchor.constraint(equalTo: thisNodeView.centerXAnchor),
                rightConnector.trailingAnchor.constraint(equalTo: rightNodeView.centerXAnchor),

                leftConnector.widthAnchor.constraint(equalTo: rightConnector.widthAnchor),

                containerView.bottomAnchor.constraint(greaterThanOrEqualTo: leftNodeView.containerView.bottomAnchor),
                containerView.bottomAnchor.constraint(greaterThanOrEqualTo: rightNodeView.containerView.bottomAnchor),
                containerView.leftAnchor.constraint(equalTo: leftNodeView.containerView.leftAnchor),
                containerView.rightAnchor.constraint(equalTo: rightNodeView.containerView.rightAnchor)
            ])
        }

        return thisNodeView
    }
}

That might look ugly, but I think it’s better than writing your own rule-based node-location engine. But the rules that these constraints capture a few basic “rules”:

1. Each node has a particular fixed size.

2. Each node has a certain distance to the level below.

3. When a node has children, center the node above the two children and space the children by a certain fixed distance.

4. When spacing peer nodes, wrap the whole binary tree below a given node in a container view, and use that for spacing. So, looking at one of the lower nodes, the - on the left side of the binary tree, the container view for its children is as follows:

   

   and when looking at the node above this, its container contains not only the two immediate children, but their containers, too:

   

   The net effect is a binary tree where all the children have reasonable spacing but parent nodes are still centered over their two immediate children.

Anyway, the view controller can call the above like so:

override func viewDidLoad() {
    super.viewDidLoad()

    // leaf nodes
    let node5 = BinaryTree.node(.empty, "5", .empty)
    let nodeA = BinaryTree.node(.empty, "a", .empty)
    let node10 = BinaryTree.node(.empty, "10", .empty)
    let node4 = BinaryTree.node(.empty, "4", .empty)
    let node3 = BinaryTree.node(.empty, "3", .empty)
    let nodeB = BinaryTree.node(.empty, "b", .empty)

    // intermediate nodes on the left
    let Aminus10 = BinaryTree.node(nodeA, "-", node10)
    let timesLeft = BinaryTree.node(node5, "*", Aminus10)

    // intermediate nodes on the right
    let minus4 = BinaryTree.node(.empty, "-", node4)
    let divide3andB = BinaryTree.node(node3, "/", nodeB)
    let timesRight = BinaryTree.node(minus4, "*", divide3andB)

    // root node
    let tree = BinaryTree.node(timesLeft, "+", timesRight)

    let nodeView = tree.addNodes(to: view)!
    NSLayoutConstraint.activate([
        nodeView.containerView.centerXAnchor.constraint(equalTo: view.centerXAnchor),
        nodeView.containerView.centerYAnchor.constraint(equalTo: view.centerYAnchor)
    ])
}

Yielding: