Skip to main content
8 of 9
credit before deletion
rolfl
  • 98.2k
  • 17
  • 220
  • 419

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));
}
rolfl
  • 98.2k
  • 17
  • 220
  • 419