Python自定义多叉树指定规则中序遍历迭代器实现排查

原代码核心逻辑错误 #

• 语法缩进错误：__next__方法被错误缩进在__iter__方法的代码块内部，不属于tree类的方法，无法被Python迭代协议正常调用
• 栈操作逻辑错误：栈是后进先出结构，原代码正序压入所有子节点，会导致最先弹出的是最右侧子节点，完全颠倒遍历顺序；同时用两个列表分别记录已访问、已返回节点的逻辑混乱，单步迭代下循环检查未返回节点的写法会直接跳过节点、重复返回节点
• 遍历边界计算错误：后半段子节点的起始索引计算不符合预期遍历规则，会重复遍历前半段节点
• 测试用例笔误：原测试代码中node4=node(3)赋值错误，会导致输出出现重复的Node3，实际应赋值为4

正确实现思路 #

放弃冗余的访问状态记录逻辑，栈中存储(节点对象, 是否已处理)二元组维护遍历状态，利用栈后进先出的特性保证遍历顺序：

1. 初始化时将根节点以未处理状态压入栈
2. 每次调用__next__时弹出栈顶元素：
   • 如果弹出的是已处理节点，直接返回该节点
   • 如果弹出的是未处理节点，按照「后遍历的内容先压栈」的规则依次压入三部分内容：
     1. 逆序压入后半段所有子节点，全部标记为未处理
     2. 压入当前节点，标记为已处理（第二次弹出时直接返回）
     3. 逆序压入前半段所有子节点，全部标记为未处理
3. 栈为空时抛出StopIteration结束遍历
   其中前半段子节点长度为子节点总数 - 子节点总数//2，剩余节点为后半段，完全匹配期望的遍历顺序：n个子节点时，先遍历前半段子节点的完整子树，再返回当前节点，最后遍历后半段子节点的完整子树。

修正后完整可运行代码 #

class node:
    def __init__(self, val=1):
        self.val = val
        self.children = []
        self.parent = None
    def add_child(self, child):
        self.children.append(child)
        child.parent = self
    def __repr__(self):
        return f"Node{self.val}"

class Stack:
    def __init__(self):
        self.data = []
        self.size = 0
    def push(self, element):
        self.size += 1
        self.data.append(element)
    def pop(self):
        res = self.data.pop(self.size - 1)
        self.size -= 1
        return res

class tree:
    def __init__(self, root):
        self.root = root
    def __iter__(self):
        self.stack = Stack()
        # 栈元素格式：(节点对象, 是否已处理标记)
        self.stack.push((self.root, False))
        return self
    def __next__(self):
        while self.stack.size > 0:
            curr_node, is_processed = self.stack.pop()
            if is_processed:
                return curr_node
            child_count = len(curr_node.children)
            # 计算前半段子节点数量
            mid = child_count - child_count // 2
            # 第一步：逆序压入后半段子节点，保证弹出顺序和children列表顺序一致
            for i in range(child_count - 1, mid - 1, -1):
                self.stack.push((curr_node.children[i], False))
            # 第二步：压入当前节点，标记为已处理
            self.stack.push((curr_node, True))
            # 第三步：逆序压入前半段子节点
            for i in range(mid - 1, -1, -1):
                self.stack.push((curr_node.children[i], False))
        raise StopIteration

if __name__ == '__main__':
    node6 = node(6)
    node3 = node(3)
    node5 = node(5)
    node9 = node(9)
    node1 = node(1)
    node2 = node(2)
    node4 = node(4) # 修正原测试用例的val赋值错误
    node8 = node(8)
    node10 = node(10)
    node7 = node(7)
    node6.add_child(node3)
    node6.add_child(node5)
    node6.add_child(node9)
    node3.add_child(node1)
    node3.add_child(node2)
    node3.add_child(node4)
    node9.add_child(node8)
    node9.add_child(node10)
    node8.add_child(node7)
    mytree = tree(node6)
    for n in mytree:
        print(n)

运行结果 #

运行上述代码后输出完全符合预期：

Node1
Node2
Node3
Node4
Node5
Node6
Node7
Node8
Node9
Node10

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