Python树的平衡检测算法实现

class TreeNode:
    def __init__(self, value):
        self.val = value
        self.left = None
        self.right = None

def is_balanced(root):
    def height(node):
        if not node:
            return 0
        left_height = height(node.left)
        right_height = height(node.right)
        return 1 + max(left_height, right_height)

    if not root:
        return True

    left_height = height(root.left)
    right_height = height(root.right)

    if abs(left_height - right_height) <= 1:
        return is_balanced(root.left) and is_balanced(root.right)
    else:
        return False

# 平衡二叉树
"""
        1
       / \
      2   3
     / \
    4   5
"""
balanced_tree = TreeNode(1)
balanced_tree.left = TreeNode(2)
balanced_tree.right = TreeNode(3)
balanced_tree.left.left = TreeNode(4)
balanced_tree.left.right = TreeNode(5)

# 非平衡二叉树
"""
        1
       / \
      2   3
         / \
        4   5
       /
      6
"""
unbalanced_tree = TreeNode(1)
unbalanced_tree.left = TreeNode(2)
unbalanced_tree.right = TreeNode(3)
unbalanced_tree.right.left = TreeNode(4)
unbalanced_tree.right.right = TreeNode(5)
unbalanced_tree.right.left.left = TreeNode(6)

result_balanced = is_balanced(balanced_tree)
print("是否为平衡二叉树:", result_balanced)

result_unbalanced = is_balanced(unbalanced_tree)
print("是否为平衡二叉树:", result_unbalanced)