Как вывести все элементы на каждом уровне глубины BST? - PullRequest
Новая
       19

Как вывести все элементы на каждом уровне глубины BST?

2 голосов
/ 24 марта 2019

Я пытаюсь создать генератор, который даст мне значения следующего уровня глубины 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__, чтобы получать значения от каждогоуровень глубины дерева?

...