如何使用栈实现非递归DFS计算以u为根的子树节点数

你提到的递推逻辑本质是后序遍历的计算顺序：必须等所有子节点的子树计数全部计算完成，才能汇总得到当前节点的子树总数，所以栈实现的核心是给每个入栈的节点加访问标记，区分「首次入栈待遍历子节点」和「子节点已全部处理待计算计数」两种状态。

具体实现步骤 #

• 初始化栈，栈中每个元素是二元组 (当前节点, 是否已访问)，初始时把根节点u压入栈，状态设为False（未访问）
• 初始化一个字典count_map，用来存储每个节点对应的子树节点数
• 循环处理栈直到栈为空：
  • 弹出栈顶元素，如果状态是False（首次弹出，还未处理子节点）：
    • 把这个节点重新压回栈，状态改为True（标记后续弹出时直接计算计数）
    • 把该节点的所有子节点逆序压入栈，状态都设为False（保证出栈顺序和子节点顺序一致，不要求子节点处理顺序的场景可以省略逆序操作）
  • 如果状态是True（所有子节点都已经处理完毕，count_map中已存储所有子节点的计数结果）：
    • 初始化当前节点计数为1（统计节点自身）
    • 遍历所有子节点，把每个子节点在count_map中的值累加进去
    • 将当前节点和对应计数存入count_map
• 最终count_map[u]就是以u为根的子树总节点数

示例代码（Python） #

class TreeNode:
    def __init__(self, val, children=None):
        self.val = val
        self.children = children if children else []

def count_subtree_nodes(root):
    if not root:
        return 0
    stack = [(root, False)]
    count_map = {}
    while stack:
        node, visited = stack.pop()
        if not visited:
            stack.append((node, True))
            # 逆序压入子节点，保证处理顺序和子节点定义顺序一致
            for child in reversed(node.children):
                stack.append((child, False))
        else:
            cnt = 1
            for child in node.children:
                cnt += count_map[child]
            count_map[node] = cnt
    return count_map[root]

适用场景说明 #

该实现完全匹配你提到的递推公式，没有依赖递归调用栈，同时严格保证了计算当前节点前，所有子节点的计数都已经完成。如果是二叉树场景，只需要把子节点的处理逻辑换成左、右子节点即可，整体逻辑完全通用。

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