Problem overview
The function isBalanced receives root, a pointer to the top node of a binary tree, 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 lesson provides the TreeNode type, with an integer Val and two pointers, Left and Right, that are nil when a child is missing.
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.
In the first example, the root 3 has the left child 9 and the right child 20, which has the children 15 and 7. At the root, the left subtree has height 1 and the right subtree has height 2: the difference is 1, which is allowed. The whole tree is balanced, and the answer is true.
In the second example, the root 1 has a left subtree of height 3 (the 2 with the 3s and 4s below it) and a right subtree of height 1: the difference is 2, so the answer is false. The third example is trickier. The root's two subtrees both have height 3, so the root is fine. But the left child 2 has a left subtree of height 2 (the 3 and the 4) and an empty right subtree of height 0, a difference of 2. The answer is false, because the rule must hold at every node, not only at the root.
The lesson asks for O(n) time, where n is the number of nodes, visiting each node once.
Approach: compute the height and check the rule together
The obvious solution has two parts: a function that measures a subtree's height, and a recursion that calls it at every node to compare the two sides. It is correct, but each call to the height function walks a whole subtree, and every node's height is measured again for each ancestor above it. On a tall, thin tree this adds up to more than O(n) work.
The faster way finds each subtree's height once, from the bottom up, and checks the rule at the same time. A recursive helper height returns the height of the tree under a node when that tree is balanced, and -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", and it passes up through the recursion. Then isBalanced only asks whether height(root) is not -1.
For the trees of the lesson, the helper gives these heights. For the first example it returns 3. For the second, the root's left subtree has height 3, the right one 1, and the helper returns -1 for the whole tree. In the third, the left child 2 returns -1 and so does the right child 2, and the signal makes the root return -1 too:
true 3
false -1
false -1
Each line shows the answer of isBalanced for one of the three trees, then the value that height returns for it.
Step by step
- Write
isBalancedas one line. It returnsheight(root) != -1. - Write the signature of
heightbelow it: a parameternodeof type*TreeNode, returningint. Go allows a function to call one that is written further down. - Base case. If
node == nil, return 0. An empty tree has height 0 and is balanced. This also stops the recursion when it steps past a leaf. - Left height. Call
height(node.Left)and store it inleftHeightwith:=. Trust the recursion: it returns the height of the left subtree, or-1if something inside it breaks the rule. - Pass the signal up. If
leftHeight == -1, return-1at once. When part of the tree is unbalanced, the whole tree is, so the right subtree doesn't even need a look. - Right height. Do the same with
node.RightandrightHeight. A missing child gives 0 thanks to the base case. - Pass the right signal up too. If
rightHeight == -1, return-1. This check is needed:-1would otherwise pass for a height, and it differs from an empty side's 0 by only 1. - Check the node itself. If one height is more than 1 above the other, return
-1. Every node runs this code on its way back up, so the rule is checked at every node. Go has no absolute value forint, so the code tests both differences, since either side may be taller. - Return the height. It is 1 for the node itself plus the larger of the two subtree heights.
The solution in Go
func isBalanced(root *TreeNode) bool {
return height(root) != -1
}
func height(node *TreeNode) int {
if node == nil {
return 0
}
leftHeight := height(node.Left)
if leftHeight == -1 {
return -1
}
rightHeight := height(node.Right)
if rightHeight == -1 {
return -1
}
if leftHeight-rightHeight > 1 || rightHeight-leftHeight > 1 {
return -1
}
return 1 + max(leftHeight, rightHeight)
}
Go details worth knowing
- Pointers to nodes. A missing child is a
nilpointer of type*TreeNode, and the code compares it withnil. Readingnode.Lefton anilnode would panic, which is why the base case comes first. - The built-in
max. Since Go 1.21maxis part of the language and works on ints. The lesson uses Go 1.22, so no import or helper is needed. - A sentinel value. Go functions can return only values of their declared type, so one
intcarries two meanings here: a height, or the signal-1. A second return value of typeboolwould also work, with(int, bool). - Early returns. Returning as soon as the left side is
-1skips the right side's work. In a tree where the problem is deep on the left, the right subtree is never visited. - Recursion depth. The tests have trees up to 800 levels deep. Go grows its stack as needed, so recursion of that depth is not a problem; a chain of 800 nodes is answered without trouble.
Complexity
Time: O(n). Every node is visited at most once, and each visit does constant work after its two recursive calls return.
Space: O(h), where h is the height of the tree: the recursion keeps one frame for each level on the path. That is O(log n) for a balanced tree and O(n) in the worst case, for a tree shaped like a chain.
Edge cases and common mistakes
- The empty tree.
isBalanced(nil)returnstrue: the base case gives height 0, which is not-1. - A single node. Both sides have height 0, so the result is
true. A node with one leaf child is balanced too (heights 1 and 0), but a chain of three nodes is not:1 -> 2 -> 3down the left side has heights 2 and 0 at the root. - Checking only the root. The third example shows why this fails: the root's subtrees both have height 3, yet the tree is unbalanced deeper down.
- Skipping the
-1checks. If the helper only compared the two heights, the signal would be treated as a height. Take a chain of three nodes going left, hanging under a root with no right child. The inner chain returns-1, and next to an empty right side (0) the difference is only 1, so the root would be reported as balanced. A version of the helper without the two early returns gives-1for the inner chain andtruefor the whole tree, which is wrong. - Measuring heights again at each node. It gives the right answer but repeats work, which is what the lesson avoids.
- Counting edges instead of nodes. The height here counts nodes, so an empty tree is 0 and a leaf is 1. The differences between sibling heights are the same either way, but the base case must match the choice.
Testing
- Root 3 with children 9 and 20, where 20 has children 15 and 7, gives
true. - Root 1 with a left subtree of height 3 and a right subtree of height 1 (the second example) gives
false. - The third example, where the imbalance is at the left 2 and not at the root, gives
false. - An empty tree gives
true. - A single node gives
true. - A node with one child gives
true; a chain of three nodes going down one side givesfalse. - A chain of 800 nodes gives
false.
