In this solution, we write code to determine if the Binary Tree is balanced or not.
For this problem, a height-balanced binary tree is defined as a binary tree in which the left and right subtrees of every node differ in height by no more than 1.
class Solution:
def is_balanced(self, root: Optional[TreeNode])->bool:
"""
Function to determine if Binary Tree is balanced.
Args:
root(TreeNode): optional TreeNode root
Returns:
bool: True if balanced else False
Examples:
>>> root = TreeNode(1)
>>> root.left = TreeNode(2)
>>> root.right = TreeNode(3)
>>> root.left.left = TreeNode(4)
>>> root.left.right = TreeNode(5)
>>> root.right.right = TreeNode(6)
>>> is_balanced(root)
True
"""
def dfs(root):
"""
Helper function for DFS in a binary tree
Args:
root: this is the root node of the binary tree
Returns:
[bool, int]: list containing a boolean value & an integer of the difference
"""
if not root: return [True, 0]
left, right = dfs(root.left), dfs(root.right)
balanced =(left[0] and right[0] and abs(left[1] - right[1]) <= 1)
height = 1 + max(left[1], right[1])
return [balanced, height]
# returns the bool value
return dfs(root)[0]
tr = TreeNode(3)
tr.left =TreeNode(9)
tr.right =TreeNode(20)
tr.right.left =TreeNode(15)
tr.right.right =TreeNode(7)
tr.right.right.left =TreeNode(8)
cls = Solution()
print("is balanced::", cls.is_balanced(tr))
The code begins with the definition of the Solution class.
The Method is_balanced, takes two arguments: self (implicitly passed) and root, which represents the root of the binary tree. The return type is bool.
The docstring explains the purpose of the method, arguments, return type, and an example of usage.
Helper Function dfs:
Inside the is_balanced method, there's a nested function dfs which stands for Depth First Search. This function is used to perform a depth-first traversal of the binary tree and calculate the height of each subtree.
The dfs function takes a root node as an argument and returns a list containing a boolean value indicating whether the subtree rooted at root is balanced and an integer representing the height of the subtree.
If root is None, indicating an empty subtree, the function returns [True, 0] . The latter shows that the subtree is balanced and has a height of 0.
Depth-First Traversal
The dfs function recursively calls itself on the left and right children of the current node (root).
It calculates the height of the left and right subtrees and checks if they are balanced by comparing their heights. If both subtrees are balanced and the difference in their heights is less than or equal to 1, the current subtree is balanced.
The is_balanced method returns the boolean value obtained by calling the dfs function on the root node and extracting the first element of the returned list indicating whether the entire binary tree is balanced or not.
Time Complexity
The time complexity primarily depends on the depth-first traversal (DFS) of the binary tree.
Each node in the binary tree is visited exactly once during the DFS traversal.
At each node, a constant number of operations are performed.
Therefore, the time complexity of the solution is O(n), where n is the number of nodes in the binary tree.
Space Complexity
The space complexity is determined by the space used in the call stack during the DFS traversal.
In the worst-case scenario, where the binary tree is completely unbalanced (i.e., a skewed tree), the depth of the call stack can be equal to the number of nodes in the tree.
Additionally, at each level of recursion, a constant amount of space is used to store the results (boolean value and height) for each node.
Therefore, the space complexity is O(n), where n is the number of nodes in the binary tree.
Alternative implementation
class Solution2:
def is_balanced(self, root: Optional[TreeNode]) -> bool:
"""
Function to determine if Binary Tree is balanced.
Args:
root (TreeNode): Optional TreeNode root of the binary tree.
Returns:
bool: True if the binary tree is balanced, otherwise False.
Example:
>>> tree = TreeNode(1)
>>> tree.left = TreeNode(2)
>>> tree.right = TreeNode(3)
>>> tree.left.left = TreeNode(4)
>>> tree.left.right = TreeNode(5)
>>> tree.right.right = TreeNode(6)
>>> Solution2().is_balanced(tree)
True
"""
# Helper function to calculate the height of a subtree rooted at 'node'
def get_height(node):
"""
Helper function to calculate the height of a subtree.
Parameters:
node (TreeNode): The root node of the subtree.
Returns:
int: The height of the subtree rooted at 'node'.
"""
if not node:
return 0
left_height = get_height(node.left)
right_height = get_height(node.right)
return max(left_height, right_height) + 1
# Helper function to recursively check if the binary tree is balanced
def is_balanced_helper(node):
"""
Helper function to recursively check if the binary tree is balanced.
Parameters:
node (TreeNode): The root node of the binary tree.
Returns:
bool: True if the binary tree rooted at 'node' is balanced, otherwise False.
"""
if not node:
return True
left_height = get_height(node.left)
right_height = get_height(node.right)
# Check if the absolute difference in heights of left and right subtrees is greater than 1
if abs(left_height - right_height) > 1:
return False
# Recursively check if both left and right subtrees are balanced
return is_balanced_helper(node.left) and is_balanced_helper(node.right)
# Call the helper function to check if the entire binary tree is balanced
return is_balanced_helper(root)
Time Complexity
The time complexity of the is_balanced function primarily depends on the traversal of the entire binary tree.
In the worst case, each node in the binary tree is visited once to calculate the height of its subtree, and then the balance condition is checked for each node resulting in O(n), where n is the number of nodes in the binary tree.
Space Complexity
The space complexity of the solution is determined by the recursive call stack during the depth-first traversal of the binary tree.
In the worst case, the depth of the call stack is equal to the height of the binary tree. Additionally, space is also used to store the height of each subtree during the traversal.
Therefore, the space complexity of the solution is O(h), where h is the height of the binary tree.
But, in the worst-case scenario of a skewed binary tree, where all nodes are on one side, the height of the tree approaches n, resulting in a space complexity of O(n).
However, in a balanced binary tree, the height h is logarithmic for the number of nodes n (h = log(n)), resulting in a space complexity of O(log(n)).
Unit Tests
As is our norm, we have unit tests to validate the correctness and robustness of our solution.
import unittest
class TestBalancedBinaryTree(unittest.TestCase):
class TreeNode:
"""TreeNode class Structure"""
def __init__(self, val=0, left =None, right=None):
self.val = val
self.left = left
self.right = right
def setUp(self) -> None:
super().setUp()
self.tr = TreeNode(3)
self.tr.left =TreeNode(9)
self.tr.right =TreeNode(20)
self.tr.right.left =TreeNode(15)
self.tr.right.right =TreeNode(7)
self.cls = Solution()
# Sample balanced tree
# 1
# / \
# 2 3
# / \
# 4 5
self.balanced_tree = TreeNode(1)
self.balanced_tree.left = TreeNode(2)
self.balanced_tree.right = TreeNode(3)
self.balanced_tree.left.left = TreeNode(4)
self.balanced_tree.left.right = TreeNode(5)
# Sample unbalanced tree
# 1
# / \
# 2 3
# / \
# 4 5
# /
# 6
self.unbalanced_tree = TreeNode(1)
self.unbalanced_tree.left = TreeNode(2)
self.unbalanced_tree.right = TreeNode(3)
self.unbalanced_tree.right.left = TreeNode(4)
self.unbalanced_tree.right.right = TreeNode(5)
self.unbalanced_tree.right.right.left = TreeNode(6)
def test_balanced_tree(self):
expected =True
self.assertEqual(expected, self.cls.is_balanced(self.tr))
def test_second_balanced_tree(self):
expected =True
self.assertEqual(expected, self.cls.is_balanced(self.balanced_tree))
def test_unbalanced_tree(self):
expected =False
self.assertFalse(expected, self.cls.is_balanced(self.unbalanced_tree))
if __name__=="__main__":
unittest.main()