iOS Interview - Leetcode 110. Balanced Binary Tree

May 13, 2024 • 3 min read • #swift, #interview, #leetcode
Primary idea: Checking max depth of left and right node from the parent node and compare them.
Time Complexity: O(n)
Space Complexity: O(n)
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 isBalanced(_ root: TreeNode?) -> Bool {
    return height(root) != -1
}

func height(_ root: TreeNode?) -> Int {
    // Base case: return when reaching the end of the branch.
    guard let root else { return 0}

    // Recursively calculate the height of the left and right nodes of the current parent node
    let left = height(root.left), right = height(root.right)

    // Ensure the the tree is a binary tree
    if left == -1 || right == -1 {
        return -1
    }

    // Check if the tree is balanced height
    if abs(left - right) > 1 {
        return -1
    }

    // Return the height for the current node
    return max(left, right) + 1
}

let root1 = TreeNode(3)
root1.left = TreeNode(9)
root1.right = TreeNode(20)
root1.right?.left = TreeNode(15)
root1.right?.right = TreeNode(7)
root1.printTree()
//┣━━ 3
//┃  ┣━━ 9
//┃  ┗━━ 20
//┃     ┣━━ 15
//┃     ┗━━ 7
print(isBalanced(root1)) // true

let root2 = TreeNode(1)
root2.left = TreeNode(2)
root2.right = TreeNode(2)
root2.left?.left = TreeNode(3)
root2.left?.right = TreeNode(3)
root2.left?.left?.left = TreeNode(4)
root2.left?.left?.right = TreeNode(4)
root1.printTree()
//┣━━ 3
//┃  ┣━━ 9
//┃  ┗━━ 20
//┃     ┣━━ 15
//┃     ┗━━ 7
print(isBalanced(root2)) // false

let root3: TreeNode? = nil
print(isBalanced(root3)) // true

// [1,null,2,null,3]
let root4 = TreeNode(1)
root4.right = TreeNode(2)
root4.right?.right = TreeNode(3)
root4.printTree()
//┣━━ 1
//┃  ┗━━ 2
//┃     ┗━━ 3
print(isBalanced(root4)) // false

// [1,2,2,3,null,null,3,4,null,null,4]
let root5 = TreeNode(1)
root5.left = TreeNode(2)
root5.right = TreeNode(2)
root5.left?.left = TreeNode(3)
root5.right?.right = TreeNode(3)
root5.left?.left?.left = TreeNode(4)
root5.right?.right?.right = TreeNode(4)
root5.printTree()
//┣━━ 1
//┃  ┣━━ 2
//┃  ┃  ┣━━ 3
//┃  ┃  ┃  ┣━━ 4
//┃  ┗━━ 2
//┃     ┗━━ 3
//┃        ┗━━ 4
print(isBalanced(root5)) // false