Swift 绘制二叉树类

DrawBinaryTree.swift

import UIKit

class TreeNode{
    let val:Int
    var left:TreeNode?
    var right:TreeNode?
    init(val:Int){
        self.val = val
    }
}

struct TreeNodeRect{
    let treeNode:TreeNode
    let rect : CGRect
    init(treeNode:TreeNode,rect:CGRect){
       self.treeNode = treeNode
        self.rect = rect
    }
    
    var center:CGPoint{
        return rect.center
    }
    
}

extension CGRect{
    init(center:CGPoint,radius:CGFloat){
        origin = CGPoint(x: center.x - radius, y: center.y - radius)
        size = CGSize(width: radius, height: radius)
    }
    
    init(center:CGPoint,width:CGFloat, height:CGFloat){
        origin = CGPoint(x: center.x - width * 0.5, y: center.y - height * 0.5)
        size = CGSize(width: width, height: height)
    }
    
    init(center:CGPoint,size:CGSize){
        origin = CGPoint(x: center.x - size.width * 0.5, y: center.y - size.height * 0.5)
        self.size = size
    }
    
    var center:CGPoint{
        return CGPoint(x: minX + width * 0.5, y: minY + height * 0.5)
    }
}


class BinaryTreeView: UIView{
    var binaryTree:TreeNode
    var radius : CGFloat = 30
    var lineSpacing : CGFloat = 60
    
    var treeNodeRects:[TreeNodeRect] = []
    
    init(frame:CGRect,binaryTree:TreeNode){
        self.binaryTree = binaryTree
        super.init(frame: frame)
        backgroundColor = UIColor.whiteColor()
        regenerateNodeList()
    }
    
   
    override func drawRect(rect: CGRect) {
        let ctx = UIGraphicsGetCurrentContext()

      
        let lineNodes = treeNodeRects
        let nodeRectByNode = {
            (node:TreeNode) -> TreeNodeRect? in
            for nr in lineNodes{
                if nr.treeNode === node{
                    return nr
                }
            }
            return nil
        }
        
        let drawLineWithPoint = {
            (from:CGPoint,to: CGPoint) -> Void in
            CGContextMoveToPoint(ctx, from.x, from.y)
            CGContextAddLineToPoint(ctx, to.x, to.y)
        }
        for nodeRect in lineNodes{
           let node = nodeRect.treeNode
            if let left = node.left{
                if let leftRect = nodeRectByNode(left){
                    drawLineWithPoint(nodeRect.center,leftRect.center)
                }
            }
            
            if let right = node.right{
                if let rightRect = nodeRectByNode(right){
                    drawLineWithPoint(nodeRect.center,rightRect.center)
                }
            }
            
        }
        
        CGContextStrokePath(ctx)
        
        UIColor(white: 0.88, alpha: 1.0)
        for nodeRect in treeNodeRects{
            CGContextFillEllipseInRect(ctx, nodeRect.rect)
        }
        
        for nodeRect in treeNodeRects{
            drawTextInRect(String(nodeRect.treeNode.val), rect: nodeRect.rect)
        }
        
    }
    
    func drawTextInRect(text:String,rect:CGRect){
        let attrs = [NSForegroundColorAttributeName: UIColor.whiteColor(),
            NSFontAttributeName: UIFont.systemFontOfSize(rect.height * 0.5)
        ]
        let attributedText = NSMutableAttributedString(string: text,attributes: attrs)
        let textSize = attributedText.size()
        let textRect = CGRect(center: CGPoint(x: rect.midX, y: rect.midY), size:textSize)
        
        attributedText.drawInRect(textRect)
        
    }
    
    private func regenerateNodeList(){
        treeNodeRects.removeAll()
        arrangeTreeNode(binaryTree, inRect: bounds)
        setNeedsDisplay()
    }
    
    func arrangeTreeNode(treeNode:TreeNode,inRect rect:CGRect){
        let nodeRect = CGRect(center: CGPoint(x: rect.midX, y: rect.minY + radius), radius: radius)
        treeNodeRects.append(TreeNodeRect(treeNode: treeNode, rect: nodeRect))
        
        let childRect = rect.rectsByDividing(nodeRect.height + lineSpacing, fromEdge: .MinYEdge).remainder
        
        let childRects = childRect.rectsByDividing(childRect.width * 0.5, fromEdge: .MinXEdge)
        if let left = treeNode.left{
           arrangeTreeNode(left, inRect: childRects.slice)
        }
        if let right = treeNode.right{
            arrangeTreeNode(right, inRect: childRects.remainder)
        }
    }

    required init(coder aDecoder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
}

let binaryTree4 = createBinaryTree4()
let binaryTreeView = BinaryTreeView(frame: CGRect(x: 0, y: 0, width: 320, height: 480), binaryTree: binaryTree4)