Returning Array from Recursive Binary Tree Search

Question

Hi I've made a simple Binary Tree and added a pre-order traversal method. After throwing around some ideas I got stuck on finding a way to return each value from the traverse_pre() method in an array.

class BST:
    def __init__(self, val):
        self.value = val
        self.left = None
        self.right = None

    def add_child(self, val):
        if self.value:
            if val < self.value:
                if self.left == None:
                    self.left = BST(val)
                else:
                    self.left.add_child(val)
            else:
                if val > self.value:
                    if self.right == None:
                        self.right = BST(val)
                    else:
                        self.right.add_child(val)
        else:
            self.value = val

    def traverse_pre(self):
        if self.left:
            self.left.traverse_pre()
        print(self.value)

        if self.right:
            self.right.traverse_pre()


Tree = BST(5)
Tree.add_child(10)
Tree.add_child(8)
Tree.add_child(2)
Tree.add_child(4)
Tree.add_child(7)

Tree.traverse_pre()

How would I modify the traverse_pre() function to return an array consisting of the node values. Is there a good example of this process for me to understand this further, I'm a bit stuck on how values can be appended to an array within recursion.

Answer 1

I would not recommend copying the entire tree to an intermediate list using .append or .extend. Instead use yield which makes your tree iterable and capable of working directly with many built-in Python functions -

class BST:
    # ...
    def preorder(self):
        # value
        yield self.value
        # left
        if self.left: yield from self.left.preorder()
        # right
        if self.right: yield from self.right.preorder()

We can simply reorder the lines this to offer different traversals like inorder -

class BST:
    # ...
    def inorder(self):
        # left
        if self.left: yield from self.left.inorder()
        # value
        yield self.value
        # right
        if self.right: yield from self.right.inorder()

And postorder -

class BST:
    # ...
    def postorder(self):
        # left
        if self.left: yield from self.left.postorder()
        # right
        if self.right: yield from self.right.postorder()
        # value
        yield self.value

Usage of generators provides inversion of control. Rather than the traversal function deciding what happens to each node, the the caller is left with the decision on what to do. If a list is indeed the desired target, simply use list -

list(mytree.preorder())

# => [ ... ]

That said, there's room for improvement with the rest of your code. There's no need to mutate nodes and tangle self context and recursive methods within your BST class directly. A functional approach with a thin class wrapper will make it easier for you to grow the functionality of your tree. For more information on this technique, see this related Q&A.

If you need to facilitate trees of significant size, a different traversal technique may be required. Just ask in the comments and someone can help you find what you are looking for.