编写二叉树高度平衡校验函数时调用LinkedBinaryTree的height()出错

现有代码的错误点 #

• 两次调用bin_tree.height()获取的都是整棵二叉树的高度，不是左、右子树的独立高度
• LinkedBinaryTree提供的height()方法默认只能计算整棵树的高度，没有入参支持指定子树根节点计算对应子树高度
• 递归调用is_height_balanced()时传入的是Node类型的子节点，不是LinkedBinaryTree实例，访问root属性时会直接报错

方案1：修改LinkedBinaryTree的height()方法适配需求 #

给height()方法增加可选的根节点入参，即可灵活计算任意子树的高度：

def height(self, calc_root=None):
    def subtree_height(root):
        if (root.left is None and root.right is None):
            return 0
        elif (root.left is None):
            return 1 + subtree_height(root.right)
        elif (root.right is None):
            return 1 + subtree_height(root.left)
        else:
            left_height = subtree_height(root.left)
            right_height = subtree_height(root.right)
            return 1 + max(left_height, right_height)

    if(self.is_empty()):
        raise Exception("Tree is empty")
    # 未指定计算根节点时默认用整棵树的根
    target_root = calc_root if calc_root is not None else self.root
    return subtree_height(target_root)

对应修改is_height_balanced函数，采用后序遍历一次完成判断，时间复杂度为O(n)，比反复调用height()的O(n²)方案效率高很多：

from LinkedBinaryTree import LinkedBinaryTree

def is_height_balanced(bin_tree):
    # 辅助函数：判断以node为根的子树是否平衡，平衡返回子树高度，不平衡返回-1
    def check_balance(node):
        if node is None:
            return 0
        left_height = check_balance(node.left)
        if left_height == -1:
            return -1
        right_height = check_balance(node.right)
        if right_height == -1:
            return -1
        if abs(left_height - right_height) > 1:
            return -1
        return max(left_height, right_height) + 1
    
    if bin_tree.root is None:
        return True
    return check_balance(bin_tree.root) != -1

方案2：不修改原LinkedBinaryTree类的实现 #

如果不允许修改原有类的代码，可以在调用时把左右子节点包装成新的LinkedBinaryTree实例再计算高度：

from LinkedBinaryTree import LinkedBinaryTree

def is_height_balanced(bin_tree):
    if bin_tree.root is None:
        return True
    # 计算左子树高度
    left_height = 0
    if bin_tree.root.left is not None:
        left_tree = LinkedBinaryTree(root=bin_tree.root.left)
        left_height = left_tree.height()
    # 计算右子树高度
    right_height = 0
    if bin_tree.root.right is not None:
        right_tree = LinkedBinaryTree(root=bin_tree.root.right)
        right_height = right_tree.height()
    
    if abs(left_height - right_height) <= 1:
        # 左右子树递归判断
        left_balanced = True
        if bin_tree.root.left is not None:
            left_balanced = is_height_balanced(LinkedBinaryTree(root=bin_tree.root.left))
        right_balanced = True
        if bin_tree.root.right is not None:
            right_balanced = is_height_balanced(LinkedBinaryTree(root=bin_tree.root.right))
        return left_balanced and right_balanced
    return False

内容的提问来源于stack exchange，提问作者xander