2
\$\begingroup\$

Implemented iterator for a binary tree and "pre" "in" and "post" order flavors. I'm looking for code review, best practices, optimizations etc.

public class IterateBinaryTree<E> {

    private TreeNode<E> root;

    /**
     * Takes in a BFS representation of a tree, and converts it into a tree.
     * here the left and right children of nodes are the (2*i + 1) and (2*i + 2)nd
     * positions respectively.
     * 
     * @param items The items to be node values.
     */
    public IterateBinaryTree(List<? extends E> items) {
        create(items);
    }

    private static class TreeNode<E> {
        TreeNode<E> left;
        E item;
        TreeNode<E> right;

        TreeNode(TreeNode<E> left, E item, TreeNode<E> right) {
            this.left = left;
            this.item = item;
            this.right = right;
        }
    }

    private void create (List<? extends E> items) {
        root = new TreeNode<E>(null, items.get(0), null);

        final Queue<TreeNode<E>> queue = new LinkedList<TreeNode<E>>();
        queue.add(root);

        final int half = items.size() / 2;

        for (int i = 0; i < half; i++) {
            if (items.get(i) != null) {
                final TreeNode<E> current = queue.poll();                
                final int left = 2 * i + 1;
                final int right = 2 * i + 2;

                if (items.get(left) != null) {
                    current.left = new TreeNode<E>(null, items.get(left), null);
                    queue.add(current.left);
                }
                if (right < items.size() && items.get(right) != null) {
                    current.right = new TreeNode<E>(null, items.get(right), null);
                    queue.add(current.right);
                }
            }
        }
    }


    /**
     * Returns the preorder representation for the given tree.
     * 
     * @return  the iterator for preorder traversal
     */
    public Iterator<E> preOrderIterator () {
        return new PreOrderItr();
    }

    private class PreOrderItr implements Iterator<E> {
        private final Stack<TreeNode<E>> stack;

        public PreOrderItr() {
            stack = new Stack<TreeNode<E>>();
            stack.add(root);
        }

        @Override
        public boolean hasNext() {
            return !stack.isEmpty();
        }

        @Override
        public E next() {
            if (!hasNext()) throw new NoSuchElementException("No more nodes remain to iterate");

            final TreeNode<E> node = stack.pop();           

            if (node.right != null) stack.push(node.right);
            if (node.left != null) stack.push(node.left);

            return node.item;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("Invalid operation for pre-order iterator.");
        }
    }

    /**
     * Returns the in-order representation for the given tree.
     * 
     * @return  the iterator for preorder traversal
     */
    public Iterator<E> inorderIterator() {
        return new InOrderItr();
    }

    private class TreeNodeDataInOrder {
        TreeNode<E> treeNode;
        boolean visitedLeftBranch;

        TreeNodeDataInOrder(TreeNode<E> treeNode, Boolean foo) {
            this.treeNode = treeNode;
            this.visitedLeftBranch = foo;
        }
    }

    private class InOrderItr implements Iterator<E> {
        private final Stack<TreeNodeDataInOrder> stack;

        public InOrderItr() {
            stack = new Stack<TreeNodeDataInOrder>();
            stack.add(new TreeNodeDataInOrder(root, false));
        }

        @Override
        public boolean hasNext() {
            return !stack.isEmpty();
        }

        @Override
        public E next() {
            if (!hasNext()) throw new NoSuchElementException("No more nodes remain to iterate");

            while (hasNext()) {
                final TreeNodeDataInOrder treeNodeData = stack.peek();
                final TreeNode<E> treeNode = treeNodeData.treeNode;

                if (!treeNodeData.visitedLeftBranch) {
                    if (treeNode.left != null) {
                        stack.add(new TreeNodeDataInOrder(treeNode.left, false));
                    }
                    treeNodeData.visitedLeftBranch = true;
                } else {
                    stack.pop();
                    if (treeNode.right != null) {
                        stack.add(new TreeNodeDataInOrder(treeNode.right, false));
                    }
                    return treeNode.item;
                }
            }
            throw new AssertionError("A node has not been returned when it should have been.");
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("Invalid operation for in-order iterator.");
        }

    }

    /**
     * Returns the post-order representation for the given tree.
     * 
     * @return  the iterator for preorder traversal
     */
    public Iterator<E> postOrderIterator() {
        return new PostOrderItr();
    }

    private class TreeNodeDataPostOrder {
        TreeNode<E> treeNode;
        boolean visitedLeftAndRightBranches;

        TreeNodeDataPostOrder(TreeNode<E> treeNode, Boolean visitedLeftAndRightBranches) {
            this.treeNode = treeNode;
            this.visitedLeftAndRightBranches = visitedLeftAndRightBranches;
        }
    }

    private class PostOrderItr implements Iterator<E> {

        private final Stack<TreeNodeDataPostOrder> stack;

        private PostOrderItr() {
            stack = new Stack<TreeNodeDataPostOrder>();
            stack.add(new TreeNodeDataPostOrder(root, false));
        }

        @Override
        public boolean hasNext() {
            return !stack.isEmpty();
        }

        @Override
        public E next() {
            if (!hasNext()) throw new NoSuchElementException("No more nodes remain to iterate");

            while (hasNext()) {
                final TreeNodeDataPostOrder treeNodeData = stack.peek();
                final TreeNode<E> treeNode = treeNodeData.treeNode;

                if (!treeNodeData.visitedLeftAndRightBranches) {
                    if (treeNode.right != null) {
                        stack.add(new TreeNodeDataPostOrder(treeNode.right, false));
                    }
                    if (treeNode.left != null) {
                        stack.add(new TreeNodeDataPostOrder(treeNode.left, false));
                    }
                    treeNodeData.visitedLeftAndRightBranches = true;
                } else {
                    stack.pop();
                    return treeNode.item;
                }
            }

            throw new AssertionError("A node has not been returned when it should have been.");
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException("Invalid operation for post-order iterator.");
        }
    }


    public static void main(String[] args) {

        Iterator<Integer> itr = null;  

        Integer[] arr = {1, 2, 3, 4, null, 6, 7};
        Integer[] arr1 = {1, 2, 3, 4, 5, 6, 7}; // full
        Integer[] arr2 = {1, 2, null, 3, null, null, null}; // left skew
        Integer[] arr3 = {1, null, 3, null, null, null, 7}; // right skew
        Integer[] arr4 = {1}; // single element


        System.out.println("---- TESTING FOR PREORDER -----------");


        IterateBinaryTree<Integer> iteratePreOrder = new IterateBinaryTree<Integer>(Arrays.asList(arr));
        System.out.print("Expected: 1 2 4 3 6 7, Actual: ");
        itr = iteratePreOrder.preOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePreOrder1 = new IterateBinaryTree<Integer>(Arrays.asList(arr1));
        System.out.print("Expected: 1 2 4 5 3 6 7, Actual: ");
        itr = iteratePreOrder1.preOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePreOrder2 = new IterateBinaryTree<Integer>(Arrays.asList(arr2));
        System.out.print("Expected: 1 2 3, Actual: ");
        itr = iteratePreOrder2.preOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePreOrder3 = new IterateBinaryTree<Integer>(Arrays.asList(arr3));
        System.out.print("Expected: 1 3 7, Actual: ");
        itr = iteratePreOrder3.preOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePreOrder4 = new IterateBinaryTree<Integer>(Arrays.asList(arr4));
        System.out.print("Expected: 1, Actual: ");
        itr = iteratePreOrder4.preOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        System.out.println("---- TESTING FOR INORDER -----------");

        IterateBinaryTree<Integer> iterateInOrder = new IterateBinaryTree<Integer>(Arrays.asList(arr));
        System.out.print("Expected: 4 2 1 6 3 7, Actual: ");
        itr = iterateInOrder.inorderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iterateInOrder1 = new IterateBinaryTree<Integer>(Arrays.asList(arr1));
        System.out.print("Expected: 4 2 5 1 6 3 7, Actual: ");
        itr = iterateInOrder1.inorderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iterateInOrder2 = new IterateBinaryTree<Integer>(Arrays.asList(arr2));
        System.out.print("Expected: 3 2 1, Actual: ");
        itr = iterateInOrder2.inorderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iterateInOrder3 = new IterateBinaryTree<Integer>(Arrays.asList(arr3));
        System.out.print("Expected: 1 3 7, Actual: ");
        itr = iterateInOrder3.inorderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iterateInOrder4 = new IterateBinaryTree<Integer>(Arrays.asList(arr4));
        System.out.print("Expected: 1, Actual: ");
        itr = iterateInOrder4.inorderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        System.out.println("---- TESTING FOR POSTORDER -----------");

        IterateBinaryTree<Integer> iteratePostOrder = new IterateBinaryTree<Integer>(Arrays.asList(arr));
        System.out.print("Expected: 4 2 6 7 3 1, Actual: ");
        itr = iteratePostOrder.postOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePostOrder1 = new IterateBinaryTree<Integer>(Arrays.asList(arr1));
        System.out.print("Expected: 4 5 2 6 7 3 1, Actual: ");
        itr = iteratePostOrder1.postOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePostOrder2 = new IterateBinaryTree<Integer>(Arrays.asList(arr2));
        System.out.print("Expected: 3 2 1, Actual: ");
        itr = iteratePostOrder2.postOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePostOrder3 = new IterateBinaryTree<Integer>(Arrays.asList(arr3));
        System.out.print("Expected: 7 3 1, Actual: ");
        itr = iteratePostOrder3.postOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }

        System.out.println();

        IterateBinaryTree<Integer> iteratePostOrder4 = new IterateBinaryTree<Integer>(Arrays.asList(arr4));
        System.out.print("Expected: 1, Actual: ");
        itr = iteratePostOrder4.postOrderIterator();
        while (itr.hasNext()) {
            System.out.print(itr.next() + " ");
        }
    }
}
\$\endgroup\$
1

2 Answers 2

3
\$\begingroup\$

I haven't checked the correctness of the implementation, just some notes about the code:

  1. Stack seems more or less deprecated. Javadoc says the following:

    A more complete and consistent set of LIFO stack operations is provided by the Deque interface and its implementations, which should be used in preference to this class.

  2. public class IterateBinaryTree<E> {
    

    Class names should be nouns, not verbs. I'd call it IterableBinaryTree instead.

  3. TreeNode(TreeNode<E> left, E item, TreeNode<E> right) {
    

    The first and last parameters are unnecessary, the code always call the constructor with null values for left and right.

  4. Boolean (in the parameters) could the primitive boolean here:

    TreeNodeDataPostOrder(TreeNode<E> treeNode, Boolean visitedLeftAndRightBranches) {
        this.treeNode = treeNode;
        this.visitedLeftAndRightBranches = visitedLeftAndRightBranches;
    }
    

    It will throw an NPE when it's called with null so prevent it with a compile time check

    Anyway, in every existing callee calls it with false, it can be simply removed:

    TreeNodeDataPostOrder(TreeNode<E> treeNode) {
        this.treeNode = treeNode;
    }
    
  5. foo could have a descriptive name:

     TreeNodeDataInOrder(TreeNode<E> treeNode, Boolean foo) {
         this.treeNode = treeNode;
         this.visitedLeftBranch = foo;
     }
    

    The issue is the same here, the code always pass it a false, so it can be removed too:

    TreeNodeDataInOrder(TreeNode<E> treeNode) {
        this.treeNode = treeNode;
    }
    
  6. for (int i = 0; i < half; i++) {
        if (items.get(i) != null) {
            final TreeNode<E> current = queue.poll();                
            final int left = 2 * i + 1;
            final int right = 2 * i + 2;
    
            if (items.get(left) != null) {
                current.left = new TreeNode<E>(items.get(left));
                queue.add(current.left);
            }
            if (right < items.size() && items.get(right) != null) {
                current.right = new TreeNode<E>(items.get(right));
                queue.add(current.right);
            }
        }
    }
    

    You can use a guard clause here to make the code flatten.

    for (int i = 0; i < half; i++) {
        if (items.get(i) == null) {
            continue;
        }
        final TreeNode<E> current = queue.poll();
        final int left = 2 * i + 1;
        ...
    }
    
  7. I'd avoid abbreviations like Itr (in PreOrderItr and in others). I've found that they undermine autocomplete. (When you type PreOrderIte and press Ctrl+Space in Eclipse it founds nothing.)

  8. private final Stack<TreeNode<E>> stack;
    
    public PreOrderItr() {
        stack = new Stack<TreeNode<E>>();
        stack.add(root);
    }
    

    You could put the object creation to the same line as the field declaration:

    private class PreOrderItr implements Iterator<E> {
        private final Stack<TreeNode<E>> stack = new Stack<TreeNode<E>>();;
    
        public PreOrderItr() {
            stack.add(root);
        }
    

    I think it's easier to read and a line shorter.

  9. if (!treeNodeData.visitedLeftBranch) {
        addStackIfNotNull(treeNode.left);
        if (treeNode.left != null) {
            stack.add(new TreeNodeDataInOrder(treeNode.left));
        }
        treeNodeData.visitedLeftBranch = true;
    } else {
        stack.pop();
        addStackIfNotNull(treeNode.right);
        if (treeNode.right != null) {
            stack.add(new TreeNodeDataInOrder(treeNode.right));
        }
        return treeNode.item;
    }
    

    You could extract out a method here:

    private void addStackIfNotNull(final TreeNode<E> node) {
        if (node != null) {
            stack.add(new TreeNodeDataInOrder(node));
        }
    }
    

    Usage:

    if (!treeNodeData.visitedLeftBranch) {
        addStackIfNotNull(treeNode.left);
        treeNodeData.visitedLeftBranch = true;
    } else {
        stack.pop();
        addStackIfNotNull(treeNode.right);
        return treeNode.item;
    }
    
\$\endgroup\$
0
\$\begingroup\$

A minor bug: You get an IndexOutOfBoundsException for an empty list in IterateBinaryTree.create(List<? extends E> items). You don't have a comment stating you need to input a list containing at least something. Consider returning IllegalArgumentException and adding a comment.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.