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

编辑：我解决了代码中的问题：用raise StopIteration替换return iter([])与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)

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

1 个答案: