HackerRank Tree traversal

Given a Binary Search Tree, give its node values following the in-order, pre-order, post-order, and level-order traversals, using both recursive and iterative algorithms. Actually, the HackerRank problem that gave me the idea for this post asks only to implement the level-order traversal algorithm, and you can write the code as you like, so probably in the iterative way, that comes more natural in this case.

Firstly, I sketched a BST on paper, to better visualize the problem. Nothing fancy, but not too basic too. Something like this.

Then I wrote the unit test for the eight functions I am about to write, using Python as implementation language. That lead me naturally to describe the data type Node that they all use, and setting up each test translating the picture I had of my BST in a data structure in my code.

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

# ...
  
class TestCodeEval(unittest.TestCase):

    def setUp(self):
        node1 = Node(1)
        node2 = Node(2, node1)
        node4 = Node(4)
        node7 = Node(7)
        node5 = Node(5, node4, node7)

        self.root = Node(3, node2, node5)

As you can see, I gave to Node the bare minimal stuff to make it work.

def test_in_order_r(self):
    self.assertEqual([1, 2, 3, 4, 5, 7], in_order_r(self.root))

def test_in_order_i(self):
    self.assertEqual([1, 2, 3, 4, 5, 7], in_order_i(self.root))

def test_pre_order_r(self):
    self.assertEqual([3, 2, 1, 5, 4, 7], pre_order_r(self.root))

def test_pre_order_i(self):
    self.assertEqual([3, 2, 1, 5, 4, 7], pre_order_i(self.root))

def test_post_order_r(self):
    self.assertEqual([1, 2, 4, 7, 5, 3], post_order_r(self.root))

def test_post_order_i(self):
    expected_values = [1, 2, 4, 7, 5, 3]
    result = post_order_i(self.root)
    self.assertEqual(len(expected_values), len(result))
    for expected, actual in zip(expected_values, result):
        self.assertEqual(expected, actual)

def test_level_order_r(self):
    self.assertEqual([3, 2, 5, 1, 4, 7], level_order_r(self.root))

def test_level_order_i(self):
    self.assertEqual([3, 2, 5, 1, 4, 7], level_order_i(self.root))

The 'r' or 'i' at the end of the function names hint which version of the algorithm I plan to implement in there, recursive or iterative.
I have written the post-order iterative test in a slightly different way to hide the fact that I let this function return a deque instead of a plain list. Just to stress that is a bit different from all the other also when it comes to write it. In production code you don't usually want to leave such a thing. Better to convert the return value inside the function.

In order

The in/pre/post order traversals are so intrinsically recursive that could be easily written as a one-liner in Python. I have force myself to be more readable, still the code is very compact and nice to see.

def in_order_r(node):
    if not node:
        return []
    return in_order_r(node.left) + [node.data] + in_order_r(node.right)

I go down to fetch the most to the left node, then up, completing the left branch before looking in the right one.

If we don't want to use recursion, we need to explicitly use a stack.

def in_order_i(node):
    if not node:
        return []

    cur = node
    result = []
    stack = []
    while True:
        if cur:
            stack.append(cur)
            cur = cur.left
        elif stack:
            cur = stack.pop()
            result.append(cur.data)
            cur = cur.right
        else:
            return result

We push the nodes in the stack, until we find a leftmost node, than we pop the last node we pushed to the stack and put its value in the result. Look for what is now the leftmost node, and so on, until we consume all the nodes.

Pre order

The recursive code is almost the same.

def pre_order_r(node):
    if not node:
        return []
    return [node.data] + pre_order_r(node.left) + pre_order_r(node.right)

In this case, first we check the left subtree, than the right one, and finally we get the one on top.

The iterative version asks again for an explicit stack where to store the nodes temporary set apart. The loop is even more natural than the in-order one.

def pre_order_i(node):
    if not node:
        return []

    result = []
    stack = [node]
    while stack:
        cur = stack.pop()
        result.append(cur.data)
        if cur.right:
            stack.append(cur.right)
        if cur.left:
            stack.append(cur.left)

    return result

Push the root on the stack before starting the loop, then consume the top of the stack and add its children.

Post order

Again, a very clean recursive implementation.

def post_order_r(node):
    if not node:
        return []
    return post_order_r(node.left) + post_order_r(node.right) + [node.data]

The iterative version is not so immediate. At least, I spent some time on it thinking how was better to proceed. The trick is looking to the pre-order implementation and change, seeing the structure similarity, and acknowledging the differences. At first sight, they look completely different. However, if you reverse the result of the post-order traversal, you could see that something is there. Beside reversing, we need also to consider the branch in different order. This means, I am pushing the result in a queue (actually a deque, the python Queue is not a good choice in this context) and not in a stack, and I reverse the order of pushing nodes in the stack.

def post_order_i(node):
    if not node:
        return []

    stack = [node]
    result = deque()
    while stack:
        cur = stack.pop()
        result.appendleft(cur.data)  # 1
        if cur.left:
            stack.append(cur.left)  # 2
        if cur.right:
            stack.append(cur.right)
    return result

1. Push to the left in the result, and not to the right as in pre-order.
2. First the left node and then the right one.

Level order

This is kind of different from the other traversals. The recursive implementation is less immediate, and it usually comes more natural to implement it in an iterative way.

Here too it comes handy using a queue (again, here a deque because of Python) but instead of the stack, not for the result. This is because we want to get the value list level by level, so the first node we push to the buffer should be the first one to be consumed.

def level_order_i(node):
    result = []
    buffer = deque([node])
    while buffer:
        cur = buffer.popleft()
        result.append(cur.data)
        if cur.left:
            buffer.append(cur.left)
        if cur.right:
            buffer.append(cur.right)
    return result

If we really want to implement it in a recursive way, it is probably a good idea to have an helper function, since we have to carry on all the nodes at the same level, to have a way to check the level below.

Let the user call it passing the root, as usual, then convert it to a list an call the helper function:

def level_order_r(node):
    return level_order_rs([node])

In the result list we store, as usual, the values of the nodes we have already found, then we populate a list with the children of the nodes on this level, and we pass them to a recursive call.

def level_order_rs(nodes):
    result = []
    below = []
    for node in nodes:
        result.append(node.data)
        if node.left:
            below.append(node.left)
        if node.right:
            below.append(node.right)
    if below:
        result += level_order_rs(below)
    return result

I pushed both the unit test and the python script to GitHub.