非递归树遍历器优化：DFS下正确返回节点深度及扩展信息

完整实现代码 #

CHILD_NODES = "CHILD_NODES"  # 定义子节点的键名，可根据实际结构修改

def get_walker(walk_by_depth=True):
    def walk(root):
        # 队列存储元组：(当前节点, 节点深度, 父节点, 兄弟索引)
        # 根节点深度为0，无父节点，兄弟索引固定为0（唯一根节点）
        queue = [(root, 0, None, 0)]
        
        while queue:
            if walk_by_depth:
                # DFS模式：从队列头部弹出节点（子节点插入头部，优先深度遍历）
                node, depth, parent, sibling_idx = queue.pop(0)
            else:
                # BFS模式：从队列头部弹出节点（子节点追加尾部，按层级遍历）
                node, depth, parent, sibling_idx = queue.pop(0)
            
            # 返回包含所有所需信息的元组
            yield node, depth, parent, sibling_idx
            
            # 处理子节点（如果存在）
            if CHILD_NODES in node:
                children = node[CHILD_NODES]
                # 为每个子节点生成上下文信息
                child_entries = [
                    (child, depth + 1, node, idx)
                    for idx, child in enumerate(children)
                ]
                if walk_by_depth:
                    # DFS：子节点插入队列头部，优先遍历深度方向
                    queue = child_entries + queue
                else:
                    # BFS：子节点追加队列尾部，按层级顺序遍历
                    queue.extend(child_entries)
    
    return walk

核心改进说明 #

1. 上下文信息随队列传递：
   不再单独维护全局深度变量或计数映射，而是将节点的深度、父节点、兄弟索引直接存入队列的元组中，从根源解决了DFS模式下深度计算错误的问题。

2. 统一BFS/DFS逻辑：

   • DFS模式：将子节点插入队列头部，确保优先遍历深度方向的节点；
   • BFS模式：将子节点追加到队列尾部，严格按层级顺序遍历。

3. 高效获取兄弟索引：
   通过enumerate遍历子节点时直接获取其在兄弟列表中的位置，无需后续调用list.index()（避免同值节点的索引错误，同时提升效率）。

原代码问题分析 #

你之前的DFS模式深度计算错误，本质是因为depth_map_of_remaining_items的计数逻辑仅适配BFS的层级特性——BFS是按层批量遍历，每层节点数量固定；而DFS是路径动态变化的遍历（深入子节点深度+1，回溯父节点深度-1），全局变量和计数映射无法精准跟踪这种动态变化。

使用示例 #

假设树结构如下：

sample_tree = {
    "name": "root",
    CHILD_NODES: [
        {
            "name": "child1",
            CHILD_NODES: [{"name": "grandchild1"}, {"name": "grandchild2"}]
        },
        {"name": "child2"}
    ]
}

DFS遍历

dfs_traverse = get_walker(walk_by_depth=True)
for item in dfs_traverse(sample_tree):
    node, depth, parent, sibling_idx = item
    parent_name = parent["name"] if parent else None
    print(f"节点: {node['name']}, 深度: {depth}, 父节点: {parent_name}, 兄弟索引: {sibling_idx}")

输出：

节点: root, 深度: 0, 父节点: None, 兄弟索引: 0
节点: child1, 深度: 1, 父节点: root, 兄弟索引: 0
节点: grandchild1, 深度: 2, 父节点: child1, 兄弟索引: 0
节点: grandchild2, 深度: 2, 父节点: child1, 兄弟索引: 1
节点: child2, 深度: 1, 父节点: root, 兄弟索引: 1

BFS遍历

bfs_traverse = get_walker(walk_by_depth=False)
for item in bfs_traverse(sample_tree):
    node, depth, parent, sibling_idx = item
    parent_name = parent["name"] if parent else None
    print(f"节点: {node['name']}, 深度: {depth}, 父节点: {parent_name}, 兄弟索引: {sibling_idx}")

输出：

节点: root, 深度: 0, 父节点: None, 兄弟索引: 0
节点: child1, 深度: 1, 父节点: root, 兄弟索引: 0
节点: child2, 深度: 1, 父节点: root, 兄弟索引: 1
节点: grandchild1, 深度: 2, 父节点: child1, 兄弟索引: 0
节点: grandchild2, 深度: 2, 父节点: child1, 兄弟索引: 1

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