Printing a Binary Tree top-down (column wise)

Question

We have a binary tree, suppose like this:

       8
     /   \
    6     10
   / \   /  \
  4   7 9    12
 / \
3   5

We have to print this binary tree in top-down manner - column wise. Note that, 8, 7 & 9 would be considered in same column. So the required output should be:

3
4
6 5
8 7 9
10
12

Currently I've implemented this problem using two maps, and one queue. Here's the method traversal(Node<T> node), which is invoked first time, passing the root node:

public void traverse(Node<T> node) {

    Queue<Node<T>> queue = new LinkedList<>();

    Map<Integer, List<Node<T>>> columnMap = new TreeMap<>();
    Map<Node<T>, Integer> map = new HashMap<>();

    queue.add(node);

    while (!queue.isEmpty()) {

        Node<T> temp = queue.poll();

        if (map.isEmpty()) {
            map.put(temp, 0);
            columnMap.put(0, new ArrayList<Node<T>>(Arrays.asList(temp)));
        }

        int column = map.get(temp);

        if (temp.left() != null) {
            queue.add(temp.left());

            if (columnMap.containsKey(column - 1)) {
                columnMap.get(column - 1).add(temp.left());
            } else {
                columnMap.put(column - 1, new ArrayList<Node<T>>(Arrays.asList(temp.left())));
            }

            map.put(temp.left(), column - 1);
        }

        if (temp.right() != null) {
            queue.add(temp.right());

            if (columnMap.containsKey(column + 1)) {
                columnMap.get(column + 1).add(temp.right());
            } else {
                columnMap.put(column + 1, new ArrayList<Node<T>>(Arrays.asList(temp.right())));
            }

            map.put(temp.right(), column + 1);
        }
    }

    for (Entry<Integer, List<Node<T>>> entry: columnMap.entrySet()) {
        System.out.println("Column - " + entry.getKey() + " : " + entry.getValue());
    }
}

The idea is, starting from root node as column number 0, the left node will have column number -1, and the right node will have column number 1. Then the right node of the left node will again have column number 0. And similarly I proceed till I've exhausted all the nodes. I've used level order traversal for traversing.

The first map - Map<Node<T>, Integer> is used to get the column number for a particular node. I've used this map, so that I can get the column number of a node in O(1) time. For example, starting with root, I put a mapping - <root, 0>, in the map. And then, put an entry in columnMap - <0, root>, specifying there is currently root at column 0. I've used Queue for the purpose of level-order traversal (Well, it can be avoided by using pre-order, post-order or in-order).

Now, I remove the root from queue, and push it's two children. For left child, I've to find the column number based on it's parent (in this case, root). So, I get the column number for root from map. And subtract 1 from it to get column of left. If the column number is already in the columnMap, I just update the existing list of node for that column, else I add a new entry.

I thought of getting it done with just a single map, but I found it more clear.

Is there some better option? Currently, the code looks a bit complex to me. May be it can be improvised? I'm open to a completely different implementation too, provided it is more clear, and space and time efficient.

Answer 1

Using recursion in this instance helps a lot to reduce code duplication. Particularly, in this case, you have different code blocks depending on whether left() or right() are null.

Using recursion you only need one null check, and this removes a big chunk of code.

The queue is only needed for recursion as well, and, it's purpose was hard to figure out... you should have documented it more carefully...

The map variable is also a variable which has a poorly documented purpose... it's hard to tell what variables are on that, and why.

All in all, your suspicions are right that there's a better way to do this.....

Algorithmically it strikes me that a single Map<List<Integer>> is probably the right structure, and that recursion is a better solution.... consider a recursive function:

recurseTraverse(final Node<Integer> node, final Map<Integer, List<Node<Integer>>> columnmap, final int column) {
    if (node == null) {
        return;
    }
    List<Node<Integer>> list = columnmap.get(column);
    if (list == null) {
        list = new ArrayList<Node<Integer>>();
        columnmap.put(column, list);
    }
    list.add(node);
    recurse(node.left(), columnmap, column - 1);
    recurse(node.right(), columnmap, column + 1);

}

Then calling that with a TreeMap<Integer, List<Node<Integer>>> you should be able to get an improved version of your result.

public void traverse(Node<Integer> root) {
    TreeMap<Integer, List<Node<Integer>>> columnMap = new TreeMap<>();

    recurseTraverse(root, columnMap, 0);

    for (Entry<Integer, List<Node<Integer>>> entry: columnMap.entrySet()) {
        System.out.println("Column - " + entry.getKey() + " : " + entry.getValue());
    }
}

---

Update: It has been pointed out by Fihop that this answer produces the incorrect results. While it correctly puts each node in to the correct column, it adds them in a depth-first order:

The accepted solution actually does not work. For example:

    1
   / \
  2   3
   \
    4
     \
      5

it gives:

2
1 4
5 3 => should be 3 5???

The right solution to this problem is to perform a breadth first traversal (traversing each row of the tree at a time), which requires a queue, and recursion is no longer the best, or even a viable, solution.

In re-working this answer, I have also actually run the code (rather than just typing it in to an answer), and I have almost 2 more years of experience on this. So, there are some other things to point out.

First up, this code can now use Java 8, and the computeIfAbsent() Map method.

Next, using static methods that are also generic methods makes the dependence on Node<Integer> become the generic Node<N>. I have still worked the example using Integers, and I have invented what I believe the Node class should be. It's not exactly the way I would write it because the method names left() and right() would probably be named differently. I have also created a rudimentary toString.

Finally, to work the traversal a helper container-class ColumnFind was created to link a node to a column, and allow the temporary storage of the column for when the next level of the tree is traversed.

It pains me that the traversal and the println calls are in the same method, I feel they should be separate, but to keep the code consistent with my earlier answer, I will leave it like that. A Java 8 function passed in to the traversal would be a better solution....

So, here's the traversal code (and it is working in ideone here: http://ideone.com/p8RWmE ).

public static final <N> void traverse(Node<N> root) {

    final TreeMap<Integer, List<Node<N>>> columnMap = new TreeMap<>();
    final Queue<ColumnFind<N>> queue = new LinkedList<>();

    queueChild(0, root, queue);

    while (!queue.isEmpty()) {
        ColumnFind<N> cf = queue.remove();
        int column = cf.column;
        Node<N> node = cf.node;
        columnMap.computeIfAbsent(column, c -> new ArrayList<>()).add(node);
        queueChild(column - 1, node.left(), queue);
        queueChild(column + 1, node.right(), queue);
    }

    for (Entry<Integer, List<Node<N>>> entry: columnMap.entrySet()) {
        System.out.println("Column - " + entry.getKey() + " : " + entry.getValue());
    }
}

private static final <N> void queueChild(int column, Node<N> node, 
                 Queue<ColumnFind<N>> queue) {
    if (node == null) {
        return;
    }
    queue.add(new ColumnFind<>(column, node));
}