Python如何统计树结构中指定父节点之后的所有节点数量

你原有的count函数只能统计当前节点所属子树的总节点数（包含节点自身），而需求里的“节点2之后的节点”包含了不在2的子树范围内的节点，所以原有逻辑无法满足。

方案1：遍历映射法（通用易实现） #

这个方法不依赖额外的节点结构，适配所有遍历顺序的“先后”定义：

• 先按照你需求的遍历顺序（示例对应的是先序遍历，规则为根节点→左子树→右子树）遍历整棵树，将所有节点按顺序存入列表
• 找到目标节点在列表中的索引位置
• 最终结果为 列表总长度 - 目标节点索引 - 1，也就是列表中排在目标节点之后的元素总个数

示例代码：

# 先序遍历整棵树，返回节点顺序列表
def preorder_traverse(root, res):
    if root is None:
        return
    res.append(root)
    preorder_traverse(root.lchild, res)
    preorder_traverse(root.rchild, res)

def count_after_node(root, target_node):
    node_list = []
    preorder_traverse(root, node_list)
    # 查找目标节点的索引
    for idx, node in enumerate(node_list):
        if node == target_node:
            return len(node_list) - idx - 1
    return 0

方案2：递归向上统计（节省空间） #

如果你的节点结构中额外保存了指向父节点的指针parent，可以不用存储全量节点，直接递归计算：

• 先统计当前节点左右子树的总节点数
• 向上遍历父节点：如果当前节点是父节点的左孩子，就加上父节点右子树的总节点数
• 直到父节点为空时停止，最终的总和就是目标结果

示例代码：

# 复用你原有统计子树节点数的逻辑
def count_subtree(node):
    if node is None:
        return 0
    return 1 + count_subtree(node.lchild) + count_subtree(node.rchild)

def count_after_node(target_node):
    total = 0
    # 先累加当前节点的所有子节点
    total += count_subtree(target_node.lchild) + count_subtree(target_node.rchild)
    current = target_node
    while current.parent is not None:
        parent = current.parent
        # 如果当前节点是父节点的左孩子，累加父节点右子树的所有节点
        if current == parent.lchild:
            total += count_subtree(parent.rchild)
        current = parent
    return total

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