Problem overview
The method isBalanced receives the root of a binary tree, a TreeNode?, and returns true when the tree is height-balanced: at every node, the heights of the left and right subtrees differ by at most 1. The height of a tree is the number of nodes on its longest path from the top down to a leaf. An empty tree has a height of 0 and is balanced. The lesson provides the TreeNode class, with an integer val and the optional children left and right.
In the first example, [3, 9, 20, nil, nil, 15, 7], the left subtree of the root (just the 9) has height 1 and the right subtree (20 with 15 and 7) has height 2. They differ by 1, and the other nodes are fine too, so the answer is true. In the second, [1, 2, 2, 3, 3, nil, nil, 4, 4], the root's left subtree has height 3 and its right subtree height 1, which differ by 2, so the answer is false. The third, [1, 2, 2, 3, nil, nil, 3, 4, nil, nil, 4], is the subtle one. The root is fine, because both of its subtrees have height 3. But the left 2 has a left subtree of height 2 and an empty right subtree of height 0, which differ by 2. The rule must hold at every node, not only at the root.
The statement asks for O(n) time, where n is the number of nodes: each node is visited once.
Approach: compute heights from the bottom up, and check as you go
The obvious plan is to write a function that measures a tree's height, then at every node compare the heights of its two subtrees. That repeats work: the height of a deep subtree is measured again for every node above it. On a tree shaped like a long chain it adds up to quadratic time.
The fix is to measure each height once, from the bottom up, and test the rule at the same moment. A recursive helper, height, returns the height of the tree under a node, with one extra meaning: it returns -1 as soon as it finds a node that breaks the rule. A real height is never negative, so -1 can't be mistaken for one. It is a signal that means "unbalanced somewhere in here".
For a node, the helper does this:
- An empty tree has height 0.
- Get the height of the left subtree. If it is
-1, pass the signal up at once. - Do the same for the right subtree.
- If the two heights differ by more than 1, this node breaks the rule: return
-1. - Otherwise, return 1 for the node itself plus the larger of the two heights.
The method then returns whether the height of the root is not -1.
On the examples, the helper gives these values: for the first tree, the whole tree has height 3, its left child 1 and its right child 2. For the second tree, the left 2 has height 3 and the right 2 has height 1, so the root returns -1. For the third, the left 2 returns -1 because its children have heights 2 and 0, so the root returns -1 too, although its own two sides look equal in size.
Step by step
- Return from
isBalanced. The method returns whetherheight(root) != -1. - Write the helper's signature.
heighttakes one parameter without an argument label,_ node: TreeNode?, and returns anInt. It goes aboveisBalancedin the class. - Handle the empty tree.
guard let node = node else { return 0 }unwraps the optional, and an empty tree has height 0. This also ends the recursion when it steps past a leaf. Below theguard,nodestands for the unwrapped node. - Find the left height with a recursive call on
node.left, into the constantleftHeight. - Pass the signal up. If
leftHeight == -1, return-1without looking at the right side. When part of the tree is unbalanced, the whole tree is. - Find the right height the same way, into
rightHeight. A missing child gives 0, thanks to the base case. - Pass the signal up again. If
rightHeight == -1, return-1. Skipping this check is a trap:-1would pass for a height, and it differs from an empty side's 0 by only 1. - Check this node. If
abs(leftHeight - rightHeight) > 1, return-1. Every node runs this on its way back up, so the rule is checked at every node, not only at the root.absmakes the order of the two sides irrelevant. - Return the height:
1 + max(leftHeight, rightHeight), the node plus its taller side.
The solution in Swift
class Solution {
func height(_ node: TreeNode?) -> Int {
guard let node = node else {
return 0
}
let leftHeight = height(node.left)
if leftHeight == -1 {
return -1
}
let rightHeight = height(node.right)
if rightHeight == -1 {
return -1
}
if abs(leftHeight - rightHeight) > 1 {
return -1
}
return 1 + max(leftHeight, rightHeight)
}
func isBalanced(_ root: TreeNode?) -> Bool {
return height(root) != -1
}
}
The Swift details worth knowing:
- Optional children.
leftandrightareTreeNode?, so a missing child isnil. The helper takes aTreeNode?as well, so it can be called onnode.leftwithout unwrapping. guard let node = node. It declares a new constant namednodethat shadows the parameter and holds the unwrapped value. When the optional isnil, theelseblock runs and must leave the method. After theguard,nodeis a plainTreeNode, and its fields are used without!.- Methods calling each other.
heightcalls itself, andisBalancedcallsheight, with noself.in front. Inside a class's method, Swift finds the other methods of the class by name. The order in which the methods are written doesn't matter either. absandmaxare global functions:abs(-3)is 3 andmax(2, 5)is 5.- Recursion depth. The statement promises at most 800 levels, which recursion handles easily. A chain of 800 left children is rejected: its height is -1, and
isBalancedreturnsfalse.
Complexity
Time: O(n). Each node is visited once, and the work at a node is a few comparisons plus the two recursive calls. Each height is found once. The early returns of -1 only make it faster.
Space: O(h), where h is the height of the tree. The recursion keeps one call per level on the stack. For a balanced tree that is about log n, and in the worst case it is n.
Edge cases and common mistakes
- Empty tree.
nilhas height 0, soisBalanced(nil)returnstrue. - Single node and a root with one child.
[1]is balanced, and so is[1, 2]: the heights are 0 and 1. Adding a grandchild on the same side,[1, 2, nil, 3], makes the heights 2 and 0, and the answer isfalse. - Checking only the root. The third example has equal-looking sides at the root and is still unbalanced. The rule applies to every node.
- Forgetting to test for -1. Without the two
-1checks, a signal from below would be treated as a height. A subtree that returns-1next to an empty side's 0 differs by only 1, so the node above would accept it: without the checks,[1, 2, nil, 3, nil, 4]givestrue, although the 2 has subtrees of height 2 and 0. - Measuring height separately at each node. It works, but the same subtree is measured again and again, which can take quadratic time on a long chain.
- Counting edges instead of nodes. The height here counts nodes, so a single node has height 1 and an empty tree 0. The base case is what ties this together.
Testing
[3, 9, 20, nil, nil, 15, 7]givestrue.[1, 2, 2, 3, 3, nil, nil, 4, 4]givesfalse.[1, 2, 2, 3, nil, nil, 3, 4, nil, nil, 4]givesfalse.- An empty tree (
nil) givestrue. [1]givestrue, and[1, 2]givestrue.[1, 2, nil, 3]givesfalse.
