CODEWITHSUNDEEP
pythongenerator
anu
anu

Reputation: 1007

Creating a generator out of a recursive function

I have created a method in a BinaryTree class to traverse the tree in zigzag manner.

def zigzag_traversal(self,l2r_flag=True,node=DEFAULT):
    """
    traverse binary tree in zigzag manner
    ALG:    pass an additional boolean variable. if variable is true call left to right, 
            if variable is false call right to left
    """
    if node == DEFAULT:
        node = self.__root
    if node is None:
        return
    if node is not None:
        print node.getData()
    if l2r_flag == True:
        self.zigzag_traversal(False,node.getRight())
        self.zigzag_traversal(False,node.getLeft())
    else:
        self.zigzag_traversal(True,node.getLeft())
        self.zigzag_traversal(True,node.getRight())

The above code is correct, logically.

Now instead of printing the data out, I am trying to turn this into a generator and yield the node so that some other method can do some other things with it.

I looked at other questions/answers on stack overflow and modified my code to yield instead of print as below.

def zigzag_traversal(self,l2r_flag=True,node=DEFAULT):
    """
    traverse binary tree in zigzag manner
    ALG:    pass an additional boolean variable. if variable is true call left to right, 
            if variable is false call right to left
    """
    if node == DEFAULT:
        node = self.__root
    if node is None:
        return
    if node is not None:
        yield node
    if l2r_flag == True:
        for node in self.zigzag_traversal(False,node.getRight()):
            yield node
        for node in self.zigzag_traversal(False,node.getLeft()):
            yield node
    else:
        for node in self.zigzag_traversal(True,node.getLeft()):
            yield node
        for node in self.zigzag_traversal(True,node.getRight()):
            yield node

But the above code does not give the same output as the code that simply prints the node values out, even though I thought they are supposed to be logically the same. So clearly my attempt to turn my method into a generator is incorrect.

My question is how to transform my original code to yield the nodes instead of print them out correctly ?

Upvotes: 1

Views: 42

Answers (1)

Kevin
Kevin

Reputation: 76264

for node in self.zigzag_traversal(False,node.getRight()):

At a glance, I suspect the problem is here. You're using node both as the name of the variable being iterated over, and as an object being used in the call to traversal. If you overshadow node in the first loop, then it won't have the right value as you enter the second loop.

Try choosing a different name for your iterating variable.

for x in self.zigzag_traversal(False,node.getRight()):
    yield x
for x in self.zigzag_traversal(False,node.getLeft()):
    yield x

Alternatively, upgrade to Python 3.X and you can use the yield from statement in lieu of creating a new name just to yield it.

yield from self.zigzag_traversal(False,node.getRight())
yield from self.zigzag_traversal(False,node.getLeft())

Upvotes: 1

Related Questions