iOS Interview - Leetcode 543. Diameter of Binary Tree

public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init() { self.val = 0; self.left = nil; self.right = nil; }
    public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
    public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
        self.val = val
        self.left = left
        self.right = right
    }
}

extension TreeNode {
    public func printTree() {
        printTreeHelper(self, "", true)
    }

    private func printTreeHelper(_ node: TreeNode?, _ prefix: String, _ isLeft: Bool) {
        guard let node = node else { return }

        let arrow = isLeft ? "┣━━" : "┗━━"
        let branch = isLeft ? "┃  " : "   "

        print(prefix + arrow + " " + "\(node.val)")

        printTreeHelper(node.left, prefix + branch, true)
        printTreeHelper(node.right, prefix + branch, false)
    }
}

var maxPath = 0

func diameterOfBinaryTree(_ root: TreeNode?) -> Int {
    maxPath = 0
    guard let root else {
        return 0
    }

    maxDepth(root)
    return maxPath
}

// Calculate the max depth of the current node. Compare and update the global result as a sum of max depth of left and right branches
func maxDepth(_ node: TreeNode) -> Int {
    var maxLeft = 0
    var maxRight = 0

    if let leftChild = node.left {
        maxLeft = maxDepth(leftChild)
    }

    if let rightChild = node.right {
        maxRight = maxDepth(rightChild)
    }

    maxPath = max(maxPath, maxLeft + maxRight)

    return max(maxLeft, maxRight) + 1
}

let root1 = TreeNode(1)
root1.left = TreeNode(2)
root1.right = TreeNode(3)
root1.left?.left = TreeNode(4)
root1.left?.right = TreeNode(5)
root1.printTree()
print(diameterOfBinaryTree(root1)) // 3

let root2 = TreeNode(1)
root2.left = TreeNode(2)
root2.printTree()
print(diameterOfBinaryTree(root2)) // 1