Description
Return whether every node's left and right subtree heights differ by at most one. An empty tree is balanced. Examples use level-order arrays.
Solution
def is_balanced(root):
heights = {None: 0}
stack = [(root, False)] if root else []
while stack:
node, ready = stack.pop()
if ready:
left, right = heights[node.left], heights[node.right]
if abs(left - right) > 1:
return False
heights[node] = 1 + max(left, right)
else:
stack.append((node, True))
if node.right:
stack.append((node.right, False))
if node.left:
stack.append((node.left, False))
return TrueExamples
Example 1
- Input
[[3,9,20,null,null,15,7]]- Output
true
Every node's two subtree heights differ by at most one.
Example 2
- Input
[[1,2,2,3,3,null,null,4,4]]- Output
false
The root's subtree heights differ by two.
Example 3
- Input
[[]]- Output
true
The empty tree is balanced.
Approach
Compute heights bottom-up with iterative postorder. As soon as any pair of child heights differs by more than one, return false; otherwise store the parent's height and continue.
Time & space
O(n) time and O(n) auxiliary space, where n is the node count.