python returning a list from a recursive method

Question

Im using a binary tree described in this book problem solving with algorithms and data structures

class BinaryTree:
    def __init__(self,rootObj):
        self.key = rootObj
        self.leftChild = None
        self.rightChild = None

There is already a preorder traversal method defined as follows.

def preorder(tree):
    if tree:
        print(tree.self.key)
        preorder(tree.getLeftChild())
        preorder(tree.getRightChild())

I just want to add a return value of the list of nodes visited. So I can do something like

for i in preorder(tree):
    etc...

Im having trouble returning a list from a recursive method. The recursion stops as soon as it hits the 'return' I've tried variations using

return [tree.self.key] + preorder()

Or

yield ...

Any ideas?

Answer 1

Are you sure you want tree.self.key and not simply tree.key when you print?

Otherwise, a solution with yield from (Python 3.3+):

def preorder(tree):
    if tree:
        yield tree
        yield from preorder(tree.getLeftChild())
        yield from preorder(tree.getRightChild())

Or with simple yield:

def preorder(tree):
    if tree:
        yield tree
        for e in preorder(tree.getLeftChild()):
            yield e
        for e in preorder(tree.getRightChild()):
            yield e

Note that using yield or yield from transforms the function into a generator function; in particular, if you want an actual list (for indexing, slicing, or displaying for instance), you'll need to explicitly create it: list(preorder(tree)).

If you have a varying number of children, it's easy to adapt:

def preorder(tree):
    if tree:
        yield tree
        for child in tree.children:
            yield from preorder(child)