Damo
Damo

Reputation: 384

python returning a list from a recursive method

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?

Upvotes: 3

Views: 1265

Answers (4)

Francis Colas
Francis Colas

Reputation: 3647

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)

Upvotes: 5

Saksham Varma
Saksham Varma

Reputation: 2130

You can simply return a concatenated list for each recursive call, which combines the current key, the left subtree and the right subtree:

def preorder(tree):
        if tree:
            return ([tree.key] + preorder(tree.getLeftChild()) +
                    preorder(tree.getRightChild()))
        return []

For an n-ary tree:

def preorder(tree):
    if tree:
        lst = [tree.key]
        for child in tree.children:
            lst.extend(preorder(child))
        return lst
    return []

Upvotes: 0

Óscar López
Óscar López

Reputation: 235994

You have to actually return a value from the recursive function (currently, it's just printing values). And you should build the list as you go, and maybe clean the code up a bit - something like this:

def preorder(tree):
    if not tree:
        return []
    # optionally, you can print the key in this line: print(self.key)
    return [tree.key] + preorder(tree.leftChild) + preorder(tree.rightChild)

Upvotes: 2

kvorobiev
kvorobiev

Reputation: 5070

You can add second argument to preorder function, something like

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

...
vn = []
preorder(tree, vn)

Upvotes: 1

Related Questions