A tree is height balanced if every node follows the equation below. The equation just says the heights of the two children are within one of each other.

abs(height(node.left) - height(node.right)) <= 1

We could try to write a recursion to check this, but that's actually really hard. So now what? The trick is to make height recursive, and "hijack" it. After all, we definitely need to compute height at each node.

1. Recursion

We already know the recursive equation for the height from a previous problem. Here it is:

height(node) = (
    hL = height(node.left)
    hR = height(node.right)
    return 1 + max(hL, hR)
)

This equation is recursive, which means it calls itself on literally every node in the tree, and finds its height. This means the recursion is already computing everything we need in order to check if the tree is height balanced. A clever idea is to "hijack" this recursion, and stick in a few lines of code so that it also computes if the tree is height balanced, like this:

height(node) = (
    hL = height(node.left)
    hR = height(node.right)
    
    # stick in some code to compute `heightBalanced`
    # ------------------------- #
    if abs(hL - hR) > 1:  # if this is true
        balaned = False   # tree is not balanced
    # ------------------------- #

    return 1 + max(hL, hR)
)

2. Base case

The recursion stops when we get to a null node. In that case, the height is 0.

height(null) = 0

3. Code

To code this up, we combine the recursion and the base case. We use a global variable called "balanced" which we update as we hijack the function. Once we see that the tree is not balanced, we'll set it to False. Here's the code:

def heightBalanced(root):
    balanced = True

    def height(node):
        # base case
        if node is None:
            return 0
        
        # recursion
        hL = height(node.left)
        hR = height(node.right)

        # hijack code
        if abs(hL - hR) > 1: 
            nonlocal balanced # access `balanced` outside of this function
            balanced = False

        return 1 + max(hL, hR)

    height(root) # computes height(root) and `balanced` too
    return balanced

Time Complexity $O(n)$. We visit every node once and do O(1) operations per node, giving O(n)*O(1) = O(n).

Space Complexity $O(n)$. height calls itself until it reaches a null node, so it might get $n$ calls deep (the tree might be a flat line). This means the Call Stack might grow up to size $n$.