CODEWITHSUNDEEP
pythonrecursiongeneratoryield
Mike Pham
Mike Pham

Reputation: 497

Combining recursion and yield for tree traversal

I'm trying to combine recursion and yield to in-order traverse a tree

This is what I'm currently having. However, when I try to traverse the tree, it seems to only traverse the root node

class Tree:
    ...
    def post_order(self, node: TreeNode):
        """Yield next node in post order from node"""
        for child in node.get_children():
            self.post_order(child)
        yield node


if __name__ == '__main__':
    root = TreeNode('root')
    depth1a = TreeNode('1a')
    depth1b = TreeNode('1b')
    root.add_children(depth1a, depth1b)
    tree = Tree(root)
    for node in tree.post_order(root):
        print(node.get_element())

When I run the code, it only prints out

root

which is the element of the first node, not what I want which is

1a
1b
root

Does anyone have an idea what I did wrong?

Thanks everyone

Upvotes: 2

Views: 760

Answers (2)

Mulan
Mulan

Reputation: 135415

Mike Pham's answer is great but I wanted to share a back-tracking approach that helps us understand how we can manually build the desired sequence using direct recursion instead of for loops. It's not a better program; it's an exercise to check your mastery of generators -

from functools import reduce

def empty ():
  yield from ()

def postorder (node, backtrack = empty(), visit = False):
  if visit:
    yield node.data
    yield from backtrack
  elif node.children:
    yield from reduce \
      ( lambda it, child: postorder (child, it)
      , node.children[::-1]
      , postorder (node, backtrack, True)
      )
  else:
    yield from postorder (node, backtrack, True)

Test it out -

class node:
  def __init__ (self, data, *children):
    self.data = data
    self.children = children

tree = \
  node("a",
    node("b",
      node("e"),
      node("f",
        node("k"))),
    node("c"),
    node("d",
      node("g"),
      node("h"),
      node("i"),
      node("j")))

print(list(postorder(tree)))
# [ 'e', 'k', 'f', 'b', 'c', 'g', 'h', 'i', 'j', 'd', 'a' ]

This might help give you an appreciation for what yield actually does for you. Here's the same program without it. Minor differences in bold -

def empty ():
  return []

def postorder (node, backtrack = empty, visit = False):
  def gen ():
    if visit:
      return [ node.data ] + backtrack()
    elif node.children:
      return reduce \
        ( lambda it, child: postorder (child, it)
        , node.children[::-1]
        , postorder (node, backtrack, True)
        ) ()
    else:
      return postorder (node, backtrack, True) ()
  return gen

def run (gen):
  return gen ()

print(run(postorder(tree)))
# [ 'e', 'k', 'f', 'b', 'c', 'g', 'h', 'i', 'j', 'd', 'a' ]

Upvotes: 0

Mike Pham
Mike Pham

Reputation: 497

Thanks to 张实唯, turns out I have to use yield from. Calling a generator function does not yield from it:

class Tree:
    ...
    def post_order(self, node: TreeNode):
        """Yield next node in post order from node"""
        for child in node.get_children():
            yield from self.post_order(child)
        yield node

Upvotes: 2

Related Questions