Balanced Binary Tree

The problem description can be found in LeetCode.

Approach

A trick is to use a magic value -1 to mark a sub-tree as unbalanced. In this way, we can avoid to return both height and a possible boolean value, which could indicate if the tree is unbalanced or not.

In any case, both solutions are valid and they have the same space and time complexities

Solution

from typing import Optional


class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        def dfs(root: Optional[TreeNode]) -> int:
            if not root:
                return 0

            left_height = dfs(root.left)
            right_height = dfs(root.right)

            diff_heights = abs(left_height - right_height)
            if left_height == -1 or right_height == -1 or diff_heights > 1:
                return -1

            return 1 + max(left_height, right_height)

        return dfs(root) != -1

from typing import Optional, Tuple


class Solution:
    def isBalanced(self, root: Optional[TreeNode]) -> bool:
        def dfs(root: Optional[TreeNode]) -> Tuple[int, bool]:
            if not root:
                return 0, True

            left_height, left_balanced = dfs(root.left)
            right_height, right_balanced = dfs(root.right)
            height = 1 + max(left_height, right_height)
            if not left_balanced or not right_balanced:
                return height, False

            diff_heights = abs(left_height - right_height)
            is_balanced = diff_heights <= 1

            return height, is_balanced

        _, is_balanced = dfs(root)
        return is_balanced

Complexity

• Time: $O(n)$
• Space: $O(h)$

Where $n$ is the number of nodes in the tree and $h$ is its height. If the tree is unbalanced, it is like having a Linked List, hence $h=n$.