如何在BST的每个深度级别上产生所有元素?

时间:2019-03-24 13:03:23

标签: python recursion binary-search-tree yield

我正在尝试创建生成器,该生成器将为我提供BST的下一个深度级别的值。假设我们的树看起来像:

BST Example

我想获得以下代码的结果:

for level in tree:
    print(level)
# 3
# (1,5)
# (2,4,7)

我正在测试yield用法的不同变化,但是没有一个能给我预期的结果。这是我结尾的代码:

class Node:

    def __init__(self, element):
        self.left = None
        self.right = None
        self.data = element

    def __iter__(self):
        yield self.data
        if self.left and self.right:
            yield from self.left.__iter__() and self.right.__iter__()
        elif self.left:
            yield from self.left.__iter__()
        elif self.right:
            yield from self.right.__iter__()
        else:
            pass

    def add_element(self, element):
        if self.data > element:
            if self.left is None:
                self.left = Node(element)
            else:
                self.left.add_element(element)
        else:
            if self.right is None:
                self.right = Node(element)
            else:
                self.right.add_element(element)

    def get_structure(self):
        return (self.left.get_structure() if self.left else None, self.data, self.right.get_structure() if self.right else None)

下面是示例:

from node import *

tree = Node(3)
tree.add_element(5)
tree.add_element(1)
tree.add_element(2)
tree.add_element(4)
tree.add_element(7)

print('Structure:',tree.get_structure())        #result: (None, 1, (None, 2, None)), 3, ((None, 4, None), 5, (None, 7, None)))
print('Iteration:', end=" ")                    #reslut: 3 5 7 / expected result: 3 (1,5) (2,4,7)
for level in tree:
    print(level, end=" ")

请告诉我如何修复我的__iter__方法以从树的每个深度级别获取值?

0 个答案:

没有答案