Implementing recursive functions for trees in python class

Question

I have created a class Tree and a class Node. Now I needed to implement preOrder, postOrder and inOrder traversals. I did it using private functions. But is there a way to do the same using only one function?

class Node:
    def __init__(self, data):
        self.left = None
        self.right = None
        self.data = data

class Tree:
    def __init__(self):
        self.root = None

    # Private helper functions
    def __insert(self, data, root):
        if data < root.data:
            if root.left is None:
                root.left = Node(data)
            else:
                self.__insert(data, root.left)
        elif data >= root.data:
            if root.right is None:
                root.right = Node(data)
            else:
                self.__insert(data, root.right)

    # Traversals
    def __preOrder(self, root):
        print root.data
        if root.left:
            self.__preOrder(root.left)
        if root.right:
            self.__preOrder(root.right)

    # Wrapper Functions 
    def insert(self, data):
        if self.root == None:
            self.root = Node(data)
        else:
            self.__insert(data, self.root)

    def preOrder(self):
        self.__preOrder(self.root)


tree = Tree()
print "Enter elements to be inserted in the tree(End with a -1): "
while True:
    elem = int(raw_input())
    if elem == -1:
        break
    tree.insert(elem)

print "Preorder traversal: "
tree.preOrder()

Here I have to use the helper functions because I don't want the user to be providing the root element explicitly.

Answer 1

Yes, you can implement all 3 types of traversal in a single function. I've turned the traversal functions into generators to make them more versatile. So instead of printing their data they are iterators that yield their data. This lets you process the data as it's yielded, or you can capture it into a list (or other collection).

In Python 2 your classes should inherit from object, otherwise you get old-style classes, which are rather limited compared to new-style classes (Python 3 only has new-style classes).

BTW, there's no need to use double underscores for the private functions (which invokes Python's name-mangling machinery), a single leading underscore is adequate.

I've also added simple __repr__ methods to the classes, which can be handy during development & debugging.

class Node(object):
    def __init__(self, data):
        self.left = None
        self.right = None
        self.data = data

    def __repr__(self):
        return repr((self.data, self.left, self.right))


class Tree(object):
    def __init__(self):
        self.root = None

    def __repr__(self):
        return repr(self.root)

    # Private helper functions
    def _insert(self, data, root):
        if data < root.data:
            if root.left is None:
                root.left = Node(data)
            else:
                self._insert(data, root.left)
        else: # data >= root.data:
            if root.right is None:
                root.right = Node(data)
            else:
                self._insert(data, root.right)

    def _traverse(self, root, mode):
        if mode == 'pre':
            yield root.data
        if root.left:
            for u in self._traverse(root.left, mode):
                yield u
        if mode == 'in':
            yield root.data
        if root.right:
            for u in self._traverse(root.right, mode):
                yield u
        if mode == 'post':
            yield root.data

    # Wrapper Functions 
    def insert(self, data):
        if self.root == None:
            self.root = Node(data)
        else:
            self._insert(data, self.root)

    def preOrder(self):
        for u in self._traverse(self.root, 'pre'):
            yield u

    def inOrder(self):
        for u in self._traverse(self.root, 'in'):
            yield u

    def postOrder(self):
        for u in self._traverse(self.root, 'post'):
            yield u

# Test

tree = Tree()

for elem in '31415926':
    tree.insert(elem)

print tree

print "Preorder traversal: "
print list(tree.preOrder())

print "InOrder Traversal: "
print list(tree.inOrder())

print "PostOrder Traversal: "
print list(tree.postOrder())

output

('3', ('1', None, ('1', None, ('2', None, None))), ('4', None, ('5', None, ('9', ('6', None, None), None))))
Preorder traversal: 
['3', '1', '1', '2', '4', '5', '9', '6']
InOrder Traversal: 
['1', '1', '2', '3', '4', '5', '6', '9']
PostOrder Traversal: 
['2', '1', '1', '6', '9', '5', '4', '3']

Here's an example of processing the data as it's yielded:

for data in tree.inOrder():
    print data

---

FWIW, this code would be a lot cleaner in Python 3 because we could use the yield from syntax instead of those for loops. So instead of

for u in self._traverse(root.left, mode):
    yield u

we could do

yield from self._traverse(root.left, mode)