"""
Given a binary tree, determine if it is height-balanced
.
A height-balanced binary tree is a binary tree in which the depth of the two subtrees of every node never differs by more than one.
"""
from typing import Optional, Tuple
from publicmatt.leet import TreeNode
def is_balanced(root: Optional[TreeNode]) -> bool:
def dfs(root) -> Tuple[bool, int]:
if root is None:
return True, 0
left, right = dfs(root.left), dfs(root.right)
height = max(left[1], right[1]) + 1
if left[0] and right[0] and abs(left[1] - right[1]) <= 1:
return True, height
return False, height
return dfs(root)[0]
Tests
from publicmatt.leet import TreeNode
from publicmatt.leet.problems.balanced_binary_tree import is_balanced
def test_empty():
tree = TreeNode.from_list([])
assert is_balanced(tree)
def test_simple_balanced():
arr = [3, 9, 20, None, None, 15, 7]
tree = TreeNode.from_list(arr)
assert is_balanced(tree)
def test_simple_unbalanced():
arr = [1, 2, 2, 3, 3, None, None, 4, 4]
tree = TreeNode.from_list(arr)
assert not is_balanced(tree)