我正在尝试创建生成器,该生成器将为我提供BST的下一个深度级别的值。假设我们的树看起来像:
我想获得以下代码的结果:
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__
方法以从树的每个深度级别获取值?