如何在Python中实现二叉树指定节点的层级查询

你现有代码存在3个核心问题导致get_level运行不符合预期：

• insert方法左右子树插入逻辑写反：二叉搜索树中小于根节点的值应插入左子树，大于根的值插入右子树，原逻辑完全颠倒，构建出的树结构和示例不符
• get_level递归逻辑错误：递归调用子节点时没有接收返回值，子树查找结果无法回传；一进入方法就对level+1导致根节点层级计算错误；没有处理节点不存在的边界情况
• in_order方法递归调用未加self前缀，运行会直接抛出名称错误

以下是修正后的完整可运行代码，层级规则对齐你给出的示例（根节点为level 0）：

class BinaryTree:
    def __init__(self, data):
        self.left = None
        self.right = None
        self.data = data
        
    def insert(self, data):
        if self.data == data:
            return
        # 小值插左子树，大值插右子树
        if self.data > data:
            if self.left:
                self.left.insert(data)
            else:
                self.left = BinaryTree(data)
        else:
            if self.right:
                self.right.insert(data)
            else:
                self.right = BinaryTree(data)
        
    def print_tree(self, level=0):
        # 带缩进打印，方便直观查看树结构
        print('  '*level + str(self.data))
        if self.left:
            self.left.print_tree(level+1)
        if self.right:
            self.right.print_tree(level+1)
            
    def get_level(self, data, level=0):
        # 当前节点匹配，直接返回对应层级
        if self.data == data:
            return level
        # 目标值更小，递归查找左子树
        if data < self.data and self.left:
            left_res = self.left.get_level(data, level + 1)
            if left_res != -1:
                return left_res
        # 目标值更大，递归查找右子树
        if data > self.data and self.right:
            right_res = self.right.get_level(data, level + 1)
            if right_res != -1:
                return right_res
        # 遍历完未找到目标节点，返回-1作为标记
        return -1
    
    def in_order(self):
        # 修正递归调用的self前缀问题
        if self.left:
            self.left.in_order()
        print(self.data, '->', end=' ')
        if self.right:
            self.right.in_order()

调用测试 #

按你给出的示例结构构建树，调用结果完全符合预期：

if __name__ == "__main__":
    # 构建示例树
    bst = BinaryTree(10)
    bst.insert(5)
    bst.insert(15)
    bst.insert(3)
    bst.insert(7)

    print(bst.get_level(10))  # 返回0
    print(bst.get_level(5))   # 返回1
    print(bst.get_level(15))  # 返回1
    print(bst.get_level(3))   # 返回2
    print(bst.get_level(7))   # 返回2
    print(bst.get_level(999)) # 返回-1，表示节点不存在

自定义调整说明 #

• 如果你的业务规则要求根节点层级从1开始计算，只需要把get_level方法的默认参数level=0改成level=1即可
• 如果不希望用-1标记节点不存在，可以自行替换为返回None或者抛出ValueError

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