Quizzlesticks
Quizzlesticks

Reputation: 23

Reading a Binary Tree using __iter__

I have a standard Binary Tree that reads like
1
/ \
2 3
/ \ / \
4 5 6 7

I need to read the binary tree by importing my tree class, creating this tree, and then using a for loop to read it like: [4, 2, 5, 1, 6, 3, 7]

I have already made the tree and my program will generate a like tree with any amount of numbers. My problem is in the method.

def iter(self):

So far I have:

def __iter__(self):
    if self.left:
        self.left.__iter__()
    yield self
    if self.right:
        self.right.__iter__()

But when I run a for loop on the tree object like:

for item in tree: print("{}: {}").format(item.name, item.height())

It only prints the first node in my try with a correct height.

My question is basically, how to print this binary tree using recursion?

Upvotes: 2

Views: 2180

Answers (3)

Jashandeep Sohi
Jashandeep Sohi

Reputation: 5153

In Python >= 3.3, one can use the yield from x syntax for recursive iterators:

def __iter__(self):
    if self.left:
        yield from self.left
    yield self
    if self.right:
        yield from self.right

Edit: Python 3.3 support confirmed.

Upvotes: 10

myaut
myaut

Reputation: 11494

Next element from the collection should be returned by __next__(): What is the interface for python iterators?

__iter__() should return an iterator, an object that defines __next__() method. yield is used in generator which is quite different concept than iterators, but you can use generator function (because it is an iterator), i.e.:

class A:
    def __init__(self, val, a=None, b=None):
        self.left = a
        self.right = b
        self.val = val

    def __iter__(self):
        def generator():
            if self.left:
                for leaf in self.left:
                    yield leaf
            yield self
            if self.right:
                for leaf in self.right:
                    yield leaf
        return generator()

for i in A(3, A(5, A(2), A(1)), A(4)):
    print(i.val)

Upvotes: 0

Thomas Baruchel
Thomas Baruchel

Reputation: 7507

Not tried, but probably rather something like:

def __iter__(self):
    if self.left:
        for k in self.left:
            yield k
    yield self
    if self.right:
        for k in self.right:
            yield k

But check it first.

Upvotes: 0

Related Questions