Problem overview
The function is_balanced receives the root of a binary tree, an Option<Box<TreeNode>>, and returns true when the tree is height-balanced. That means that 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, and an empty tree has height 0 and counts as balanced. The lesson provides the TreeNode type, with the fields val, left and right; the two children are Option<Box<TreeNode>> values of their own.
The tree [3, 9, 20, None, None, 15, 7] is balanced. At the root the left subtree (just the 9) has height 1 and the right subtree (20 with 15 and 7) has height 2, a difference of 1, and the other nodes are fine too. The tree [1, 2, 2, 3, 3, None, None, 4, 4] is not: at the root the left side has height 3 and the right side height 1.
The third example is the one to study. In [1, 2, 2, 3, None, None, 3, 4, None, None, 4] the root is fine, because both sides have height 3. But the left 2 has a left subtree of height 2 and an empty right subtree of height 0, so the tree is unbalanced. The rule must hold at every node, not only at the root.
The lesson asks for O(n) time, visiting each node once.
Approach: heights from the bottom up, with -1 as an alarm
The obvious way writes a height function and, at every node, calls it on both children and compares. That measures the same subtrees again and again, once for every node above them, and it is slow on tall trees.
The better way finds each subtree's height once, from the bottom up, and checks the rule at the same time. A recursive helper, height, returns one of two things:
- the real height of the tree under a node, when that tree is balanced;
- the value
-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". Each node does three things on its way back up: ask the left side, ask the right side, then compare. Any -1 is passed straight up, and the final answer is just whether the root's result is not -1.
Running the helper on the examples gives these results. The first tree returns the height 3, and the second returns -1. In the third, the left 2 returns -1 because its sides have heights 2 and 0, and the root passes that on without looking at its right side.
Step by step
- Write the answer first.
is_balancedends withSelf::height(root) != -1, without a semicolon, so that is its value. - Declare the helper
heightwith a parameternode: Option<Box<TreeNode>>that returns ani32. - Make the body a
matchonnode. There is no semicolon after its closing brace, so its value is what the helper returns. - Base case: the arm
Nonegives 0. An empty tree has height 0 and is balanced. This arm also stops the recursion past a leaf. - The node arm:
Some(node)binds the box asnode, with a block body for the next steps. - Left height: call
Self::height(node.left)and keep the result inleft_height. - Pass the alarm up: if
left_height == -1,return -1;at once. The whole tree is unbalanced, so there is no need to look at the right side. - Right height: do the same with
node.rightandright_height. A missing child gives 0 thanks to the base case. - Check the right result for -1 too. Don't skip this: -1 would otherwise pass for a height, and it is only one away from an empty side's 0.
- Check this node: if
(left_height - right_height).abs() > 1, return -1. This runs at every node on the way back up, so the rule is checked everywhere. - Return the height: the arm ends with
1 + left_height.max(right_height), the node itself plus its taller side.
The solution in Rust
impl Solution {
fn height(node: Option<Box<TreeNode>>) -> i32 {
match node {
None => 0,
Some(node) => {
let left_height = Self::height(node.left);
if left_height == -1 {
return -1;
}
let right_height = Self::height(node.right);
if right_height == -1 {
return -1;
}
if (left_height - right_height).abs() > 1 {
return -1;
}
1 + left_height.max(right_height)
}
}
}
fn is_balanced(root: Option<Box<TreeNode>>) -> bool {
Self::height(root) != -1
}
}
The Rust details worth a closer look:
Option<Box<TreeNode>>. A child is eitherNoneorSomeof a boxed node. TheBoxputs the node on the heap, which a type that contains itself needs.matchmust cover both cases, and the compiler checks that it does.- Moving the subtrees.
Self::height(node.left)moves the left subtree into the call, and the same goes fornode.right. That's fine because nothing needs them afterwards, and it saves cloning. Likewiseis_balancedhandsrootto the helper. - The binding
Some(node)reuses the namenode. Inside the arm it is theBox<TreeNode>, and the fields are read through it asnode.leftandnode.right. returninside a match arm. It leaves the whole function, not just the arm. The final expression1 + left_height.max(right_height)has no semicolon, so it is the arm's value, and thematchis the function's value..abs()and.max(). Both are methods oni32.absdrops the sign, so it doesn't matter which side is taller, anda.max(b)returns the larger of two numbers.- Associated functions.
heighthas noself, so it is called asSelf::height.
Complexity
Time: O(n). Each node is visited once and does constant work besides its two recursive calls. A node can even be skipped: after an early -1, the right side is not visited.
Space: O(h), where h is the height of the tree, for the recursion. A balanced tree has h near log n, and a chain-shaped tree has h = n. The lesson's trees are at most 800 levels deep, which the recursion handles.
Edge cases and common mistakes
- An empty tree.
Nonereturns height 0, sois_balanced(None)istrue. - A single node. Both sides have height 0, so
[1]is balanced. - Short chains. A node with one child, which has no children, has side heights 1 and 0, so it is balanced. Add one more level on the same side and the heights are 2 and 0, so it is not. A chain of 800 nodes returns
false. - Checking only the root. The third example has equal heights at the root and is still unbalanced. The check must run at every node.
- Forgetting the right-hand -1 check. A -1 from the right side would be treated as a height of -1 and could pass the difference test.
- Mixing up the two meanings. A height of 0 means an empty tree, which is a fine result. Only -1 means failure.
- Using the difference without
abs. Then a taller right side would produce a negative number that never exceeds 1.
Testing
[3, 9, 20, None, None, 15, 7]givestrue.[1, 2, 2, 3, 3, None, None, 4, 4]givesfalse.[1, 2, 2, 3, None, None, 3, 4, None, None, 4]givesfalse.[1]givestrue.- An empty tree gives
true. - A root with one child, which has no children, gives
true. - A chain of three nodes, each with one left child, gives
false.
