主要涉及一个遍历层级,以及遍历计算坐标
上码
结点
class NaryTreeNode {
var parent: NaryTreeNode?
var title: String?
var son: [NaryTreeNode]?
var level: Int = 0 // 层
var frame: CGRect?
}
树配置
struct NaryTreeConfig {
var lineSpace: CGFloat = 30
var interspace: CGFloat = 30
var nodeSize = CGSize(width: 60, height: 60)
}
树
class NaryTree: UIView {
var config = NaryTreeConfig() {
didSet {
setupNodeFrame(node: root)
}
}
var root: NaryTreeNode? {
didSet {
setupNodeFrame(node: root)
}
}
override func draw(_ rect: CGRect) {
super.draw(rect)
drawNode()
}
}
private extension NaryTree {
func drawNode() {
guard let root = root else { return }
var stack = [NaryTreeNode]()
stack.append(root)
drawCircle(frame: root.frame)
while !stack.isEmpty {
let node = stack.removeLast()
guard let son = node.son else { return }
for child in son {
stack.insert(child, at: 0)
drawCircle(frame: child.frame)
}
}
}
func drawCircle(frame: CGRect?) {
guard let context = UIGraphicsGetCurrentContext(),
let frame = frame
else {
return
}
UIColor.red.set()
let path = UIBezierPath(arcCenter: CGPoint(x: frame.midX, y: frame.midY), radius: config.nodeSize.width / 2, startAngle: 0, endAngle: 360, clockwise: true).cgPath
context.addPath(path)
context.fillPath()
}
}
private extension NaryTree {
func setupNodes() {
setupNodeFrame(node: root)
setNeedsDisplay()
}
func setupNodeFrame(node: NaryTreeNode?) {
var minX: CGFloat = -config.nodeSize.width - config.interspace
guard let node = node else { return }
levelOrder(node)
calculateFrame(node: node, minX: &minX)
}
func calculateFrame(node: NaryTreeNode, minX: inout CGFloat) {
guard let son = node.son,
son.isEmpty == false
else {
minX += config.nodeSize.width + config.interspace
let y = CGFloat(node.level) * config.lineSpace + CGFloat(node.level) * config.nodeSize.height
node.frame = CGRect(x: minX, y: y, width: config.nodeSize.width, height: config.nodeSize.height)
return
}
son.enumerated().forEach { _, child in
calculateFrame(node: child, minX: &minX)
}
guard let minx = son.first?.frame?.origin.x,
let maxX = son.last?.frame?.origin.x else { return }
let x = (minx + maxX) / 2.0
let y = CGFloat(node.level) * config.lineSpace + CGFloat(node.level) * config.nodeSize.height
node.frame = CGRect(x: x, y: y, width: config.nodeSize.width, height: config.nodeSize.height)
}
/// 层序遍历
func levelOrder(_ root: NaryTreeNode?) {
guard let root = root else { return }
var stack = [NaryTreeNode]()
stack.append(root)
while !stack.isEmpty {
let node = stack.removeLast()
guard let son = node.son else { return }
for child in son {
stack.insert(child, at: 0)
child.level = node.level + 1
}
}
}
}