iOS Interview - Leetcode 235. Lowest Common Ancestor of a Binary Search 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)
    }
}

func lowestCommonAncestor(_ root: TreeNode?, _ p: TreeNode?, _ q: TreeNode?) -> TreeNode? {
    // Stop when reach end of tree
    guard let root = root else { return nil }

    // Stop and return when reaching one of the targeting node
    if root.val == p?.val || root.val == q?.val { return root }

    // Recursively check children of current root to see if one of them is the common accestor of p and q
    let left = lowestCommonAncestor(root.left, p, q)
    let right = lowestCommonAncestor(root.right, p, q)

    // If we reach both targeting node, then the current root is the common accestor
    if left != nil, right != nil { return root }

    // If the right branch doesn't reach the targeted nodes, then the left branch is the common accestor
    if left != nil { return left}

    // If the left branch doesn't reach the targeted nodes, then the right branch is the common accestor
    if right != nil { return right }

    // Worse case, there is no common accestor
    return nil
}


let root1 = TreeNode(6)
let node11 = TreeNode(2)
let node12 = TreeNode(8)
root1.left = node11
root1.right = node12
let node21 = TreeNode(0)
let node22 = TreeNode(4)
let node23 = TreeNode(7)
let node24 = TreeNode(9)
node11.left = node21
node11.right = node22
node12.left = node23
node12.right = node24
let node31 = TreeNode(3)
let node32 = TreeNode(5)
node22.left = node31
node22.right = node32

root1.printTree()
//┣━━ 6
//┃  ┣━━ 2
//┃  ┃  ┣━━ 0
//┃  ┃  ┗━━ 4
//┃  ┃     ┣━━ 3
//┃  ┃     ┗━━ 5
//┃  ┗━━ 8
//┃     ┣━━ 7
//┃     ┗━━ 9
print(lowestCommonAncestor(root1, node11, node12)?.val) // 6
print(lowestCommonAncestor(root1, node21, node32)?.val) // 2

let root2 = TreeNode(2)
let node221 = TreeNode(1)
root2.left = node221
root2.printTree()
//┣━━ 2
//┃  ┣━━ 1
print(lowestCommonAncestor(root2, root2, node221)?.val) // 2