DAG遍历需求：实现父节点优先层级遍历并修正代码问题

问题描述 #

有一个DAG（有向无环图），需要从根节点开始遍历，输出按层级分组的结果，预期格式如下：

1 6
2 3
4
5

尝试了一段Python代码但未能实现预期效果，需要明确这类遍历的算法名称、思路，并给出代码修正建议。

原尝试代码：

from typing import List


class Node:
    def __init__(self, value, children) -> None:
        super().__init__()
        self.value = value
        self.children = children

    @staticmethod
    def create(value, *parents):
        node = Node(value, [])
        if parents is not None:
            for parent in parents:
                parent.children.append(node)
        return node


def travel(roots: List[Node], visited: List[Node]):
    print(" ".join([str(r.value) for r in roots]))
    visited += roots
    all_children = []
    for r in roots:
        if r.children:
            for c in r.children:
                if c not in visited:
                    all_children.append(c)
    if all_children:
        travel(all_children, visited)


if __name__ == '__main__':
    root = Node.create(1)
    root2 = Node.create(6)
    roots = [root, root2]
    n2 = Node.create(2, root)
    n3 = Node.create(3, root)
    n4 = Node.create(4, n2, n3, root2)
    n5 = Node.create(5, n4)
    travel(roots, [])

算法说明 #

你需要的是分层拓扑遍历（也叫广度优先拓扑排序），核心思路：

• 基于拓扑排序的入度思想，跟踪每个节点的入度（即指向该节点的父节点数量）
• 只有当一个节点的所有父节点都被处理完毕（入度减至0），才将其加入当前层级
• 按层级分组输出，同一层级的节点在同一行打印

原代码的问题：仅通过visited判断是否加入下一层，导致有多个父节点的节点（比如节点4）会被提前加入队列（当第一个父节点6被处理时就会被加入），而不是等所有父节点（2、3、6）都处理完才进入下一层。

代码修正建议 #

以下是修正后的代码，采用入度跟踪的分层拓扑遍历：

from typing import List, Dict
from collections import deque


class Node:
    def __init__(self, value, children=None) -> None:
        self.value = value
        self.children = children if children is not None else []

    @staticmethod
    def create(value, *parents):
        node = Node(value)
        if parents:
            for parent in parents:
                parent.children.append(node)
        return node


def layered_topological_traversal(roots: List[Node]):
    # 计算每个节点的入度
    in_degree: Dict[Node, int] = {}
    
    def init_in_degree(node: Node):
        if node in in_degree:
            return
        in_degree[node] = 0
        for child in node.children:
            in_degree[child] = in_degree.get(child, 0) + 1
            init_in_degree(child)
    
    for root in roots:
        init_in_degree(root)
    
    # 初始化队列，根节点入度设为0（无父节点）
    queue = deque()
    for root in roots:
        in_degree[root] = 0
        queue.append(root)
    
    while queue:
        # 处理当前层级的所有节点
        level_size = len(queue)
        current_level = []
        for _ in range(level_size):
            node = queue.popleft()
            current_level.append(str(node.value))
            # 遍历子节点，入度减1
            for child in node.children:
                in_degree[child] -= 1
                if in_degree[child] == 0:
                    queue.append(child)
        # 打印当前层级
        print(" ".join(current_level))


if __name__ == '__main__':
    root = Node.create(1)
    root2 = Node.create(6)
    roots = [root, root2]
    n2 = Node.create(2, root)
    n3 = Node.create(3, root)
    n4 = Node.create(4, n2, n3, root2)
    n5 = Node.create(5, n4)
    layered_topological_traversal(roots)

代码说明 #

1. 入度计算：通过递归遍历所有节点，统计每个节点的父节点数量（入度）
2. 队列初始化：将根节点的入度设为0并加入队列
3. 分层处理：每次处理队列中当前所有节点（即同一层级），遍历它们的子节点并将子节点入度减1，当子节点入度为0时，说明所有父节点已处理完毕，加入下一层队列
4. 层级输出：每处理完一层节点，将该层节点的值拼接成字符串打印

运行这段代码后，输出将完全符合预期：

1 6
2 3
4
5

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