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)

Answer 2

I'm not sure about implementing the traversal functions as one-liners, but an alternative approach to what you're trying to do would be to adapt the Strategy Pattern to your use case by abstracting the traversal logic into a series of separate classes that all inherit from one common TraversalStrategy. You can then inject a traversal strategy object as a dependency to Tree which decouples the structure of the tree from the logic used to traverse it.

This approach has the following benefits:

• It allows you to get rid of the bulk of your current Tree methods
• It increases the testability of your program (you can unit test tree traversal strategies in isolation from other tree-related logic)
• It increases reusability (you can potentially use these traversal strategies in other situations)
• In general, it makes your program more SOLID:
  • Your Tree class now has a Single Responsibility - a single reason to change
  • You can implement new traversal behaviours for trees without having to touch the source code for Tree i.e. it's closed for modification
  • You're injecting the traversal logic as a dependency, so it's very easy to substitute between different implementations, including mocks

The code below has been written for Python 3, so it will need some minor changes to work for Python 2.

from abc import ABC, abstractmethod


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

    def __repr__(self):
        return str(self.data)


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

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

    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)

    def traverse(self):
        self.traversal_strategy.traverse(self.root)


class TraversalStrategy(ABC):
    @abstractmethod
    def traverse(self, node):
        pass

    def _attempt_traverse(self, node):
        if node:
            self.traverse(node)


class PreOrderTraversal(TraversalStrategy):
    def traverse(self, node):
        print(node)
        self._attempt_traverse(node.left)
        self._attempt_traverse(node.right)


class InOrderTraversal(TraversalStrategy):
    def traverse(self, node):
        self._attempt_traverse(node.left)
        print(node)
        self._attempt_traverse(node.right)


class PostOrderTraversal(TraversalStrategy):
    def traverse(self, node):
        self._attempt_traverse(node.left)
        self._attempt_traverse(node.right)
        print(node)


def build_tree(traversal_strategy):
    tree = Tree(traversal_strategy)
    elements = [1, 3, 6, 9, 2, 8]

    for element in elements:
        tree.insert(element)

    return tree


if __name__ == '__main__':
    pre_order_tree = build_tree(PreOrderTraversal())
    in_order_tree = build_tree(InOrderTraversal())
    post_order_tree = build_tree(PostOrderTraversal())

    print('Pre order traversal: ')
    pre_order_tree.traverse()
    print()

    print('In order traversal: ')
    in_order_tree.traverse()
    print()

    print('Post order traversal: ')
    post_order_tree.traverse()
    print()

Output

Pre order traversal: 
1
3
2
6
9
8

In order traversal: 
1
2
3
6
8
9

Post order traversal: 
2
8
9
6
3
1