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