如何使用 itertools.chain 实现二叉树遍历迭代器?

Question

我正在尝试在Python 3.7中实现一个二叉树类，但我的遍历方法无法正常工作。

我宁愿只有一个方法并选择带有参数的顺序，并且我希望它返回一个迭代器，所以我尝试递归使用 itertools.chain。 以下是我的课程：

class TreeNode:
def __init__(self, name, data=None, parent=None, mode='lcrs'):  
self.name = name
self.data = data
self.left = None
self.right = None
self.parent = None
if parent is None:
pass
else:
self.insert(parent, mode)
[...]
def traversal(self, order):
return TreeTraversal(self, order)

class TreeTraversal:
def __init__(self, root, order):
"""Iterator for the traversal identified by order
:param root: root of the tree to traverse
:param order: ordering
'pre': pre-order
'-pre': right to left pre-order
'in': in-order
'out': out-order
'post': post-order
'-post': right to left post-order
"""
par_map = dict(zip(['pre', '-pre', 'in', 'post', 'out', '-post'],
permutations(range(3))))
self.order_name = order
self.order = par_map[order]
self.root = root
def __iter__(self):
if self.root is None:
raise StopIteration
else:
call_map = dict(zip(self.order,
[[self.root],
TreeTraversal(self.root.left,
self.order_name),
TreeTraversal(self.root.right,
self.order_name)]))
call_list = [call_map[i] for i in range(3)]
return chain.from_iterable(call_list)

我试图测试这棵树：

F
/ 
B   G
/
L

L1

使用此代码：

from tree import TreeTraversal
from tree import TreeNode
F = TreeNode('F')
G = TreeNode('G', parent=F, mode='right')
B = TreeNode('B', parent=F, mode='left')
L = TreeNode('L', parent=G, mode='left')
L1 = TreeNode('L1', parent=G, mode='right')
d = {}
for o in ['pre', '-pre', 'in', 'post', 'out', '-post']:
d[o] = []
for i in F.traversal(o):
d[o].append(i.name)

我期待：

d = {'pre': ['F', 'B', 'G', 'L','L1'], 
'-pre': ['F','G','L','L1','B'], 
'in': ['B', 'F', 'L', 'L1', 'G'], 
'post': ['B', 'L1', 'L', 'G', 'F'], 
'out': ['G', 'L1', 'L', 'F', 'B'], 
'-post': ['L1', 'L', 'G', 'B', 'F']}

但相反，我得到了：

d = {'pre': ['F', 'B', 'G', 'L'], 
'-pre': ['F', 'G', 'B'], 
'in': ['F', 'G'], 
'post': ['F'], 
'out': ['F'], 
'-post': ['F']}

当其中一个子树为空时，似乎有些东西无法按预期工作，但这是我第一次尝试在教程示例之外使用迭代器，所以我真的不知道出了什么问题。

Answer 1

编辑：我解决了代码中的问题：用return iter([])替换raise StopIteration可以正确使用chain;我还意识到后订单和外订单被交换了，所以我纠正了它。 下面是迭代器类的正确代码：

class TreeTraversal:
def __init__(self, root, order):
"""Iterator for the traversal identified by order
:param root: root of the tree to traverse
:param order: ordering
'pre': pre-order
'-pre': right to left pre-order
'in': in-order
'out': out-order
'post': post-order
'-post': right to left post-order
"""
par_map = dict(zip(['pre', '-pre', 'in', 'out', 'post', '-post'],
permutations(range(3))))
self.order_name = order
self.order = par_map[order]
self.root = root
def __iter__(self):
if self.root is None:
return iter([])
else:
call_map = dict(zip(self.order,
[[self.root],
TreeTraversal(self.root.left,
self.order_name),
TreeTraversal(self.root.right,
self.order_name)]))
call_list = [call_map[i] for i in range(3)]
return chain.from_iterable(call_list)