Tree Iterator in Java

I work on a CMS during daylight hours. We model the destinations of our content as objects called Pages.

I was writing code to import a set of Pages, each of which may have a parent Page. I needed to group them in the correct order, so then when I create a Page, its parent Page has already been created.

I tried some weird stuff with sorting and Comparators. It sucked, and it wasn’t working. Then a colleague suggested a Tree. Why not?

I took a long look at TreeMap, but couldn’t figure out how to make it work for my situation.

Here is how I define a TreeNode:

public interface TreeNode<T> {

    void add(TreeNode<T>... children);

    TreeNode<T> add(T child);

    List<TreeNode<T>> add(T... children);

    T value();

    List<TreeNode<T>> childNodes();

    List<T> childValues();

    Iterator<TreeNode<T>> iterator();

    Iterator<T> valueIterator();

}

I created a simple implementation called LinkedTreeNode that stores the node’s children in a LinkedList.

This is a slightly non-standard tree, since my root doesn’t have a value. There is no root Page, just a group of top-level Pages without parents.

Here is how I built the Tree. It’s ugly (the generics don’t help), but I wanted to illustrate how I churn through the list of Pages.

TreeNode<Page> createTree(List<Page> pages) {
    final TreeNode<Page> tree = new LinkedTreeNode<Page>(null);

    // finds top-level pages
    for (final Iterator<Page> i = pages.iterator(); i.hasNext();) {
        final Page page = i.next();
        final String parentId = page.getParentId();
        if (parentId == null) {
            tree.add(new TreeNode<Page>(page));
            // removes top-level pages from list
            i.remove();
        }
    }

    for (final TreeNode<Page> parent : tree.children()) {
        // find children for page
        addChildren(parent, pages);
    }

    return tree;
}

private void addChildren(TreeNode<Page> parent,
    List<Page> pages) {
    for (final Page childPage :
        removeChildren(parent.value().getId(), pages)) {
        final TreeNode<Page> childNode =
            new TreeNode<Page>(childPage);
        parent.add(childNode);
        addChildren(childNode, pages);
    }
}

private List<Page> removeChildren(String parentId,
    List<Page> pages) {
    final List<Page> children = new ArrayList<Page>();
    for (final Iterator<Page> i = pages.iterator(); i.hasNext();) {
        final Page page = i.next();
        if (page.getParentId().equals(parentId)) {
            children.add(page);
            // removes child from list
            i.remove();
        }
    }
    return children;
}

The purpose of building this Tree is so I could iterate through the Pages in order. From each TreeNode’s perspective, I wanted to return the node, then the node’s children. I implemented it as a very special iterator, an inner class of LinkedTreeNode.

private final class DepthFirstNodeIterator
    implements Iterator<TreeNode<T>> {

    private TreeNode<T> currentNode = LinkedTreeNode.this;

    private final Iterator<TreeNode<T>> childIterator =
        children.iterator();

    private Iterator<TreeNode<T>> currentIterator;

    private final boolean hasChildren;

    public DepthFirstNodeIterator() {
        if (this.childIterator.hasNext()) {
            this.currentIterator =
                this.childIterator.next().iterator();
            this.hasChildren = true;
        } else {
            this.hasChildren = false;
        }
    }

    public boolean hasNext() {
        return this.currentNode != null;
    }

    public TreeNode<T> next() {
        final TreeNode<T> result = this.currentNode;

        if (this.hasChildren) {
            if (this.currentIterator.hasNext()) {
                this.currentNode = this.currentIterator.next();
            } else {
                if (this.childIterator.hasNext()) {
                    this.currentIterator =
                        this.childIterator.next().iterator();
                    this.currentNode =
                        this.currentIterator.next();
                } else {
                    this.currentNode = null;
                }
            }
        } else {
            this.currentNode = null;
        }

        return result;
    }

    public void remove() {
        throw new UnsupportedOperationException();
    }

}

I wrote some tests to verify that my tree works as expected.

public final class LinkedTreeNodeTest {

    @Test
    public void testValue() {
        final String value = "one";
        final TreeNode<String> tree =
            new LinkedTreeNode<String>(value);
        assertEquals(tree.value(), value, "Unexpected value.");
    }

    @Test
    public void testChildren() {
        final String[] children = {"two", "three", "four"};
        final TreeNode<String> tree =
            new LinkedTreeNode<String>("one");
        for (final String child : children) {
            tree.add(child);
        }

        final List<String> childValues = tree.childValues();
        for (int i = 0; i < childValues.size(); i++) {
            assertEquals(childValues.get(i), children[i],
                    "Unexpected child value.");
        }
    }

    @Test
    public void testIterator() {
        final TreeNode<String> tree =
            new LinkedTreeNode<String>("one");
        final TreeNode<String> two = tree.add("two");
        final TreeNode<String> three = tree.add("three");
        two.add("four", "five");
        three.add("six", "seven");

        final List<String> values = new ArrayList<String>();
        for (final Iterator<String> i = tree.valueIterator();
            i.hasNext();) {
            values.add(i.next());
        }

        final String msg = "Unexpected value.";
        assertEquals(values.get(0), "one", msg);
        assertEquals(values.get(1), "two", msg);
        assertEquals(values.get(2), "four", msg);
        assertEquals(values.get(3), "five", msg);
        assertEquals(values.get(4), "three", msg);
        assertEquals(values.get(5), "six", msg);
        assertEquals(values.get(6), "seven", msg);
    }

}