LeetCode 173 Binary Search Tree Iterator

Leetcode Binary Search Tree Iterator

/**
 * Definition for a binary tree node.
 * Basic building block for Binary Search Tree implementation.
 */
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode() {}

    TreeNode(int val) {
        this.val = val;
    }

    TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

PostOrder

import java.util.Stack;

/**
 * Iterator for Binary Search Tree that traverses in postorder (Left -> Right -> Root)
 * Uses O(h) space where h is the height of the tree
 */
public class PostorderIterator {
    // Stack to keep track of nodes in the path from root to current node
    private Stack<TreeNode> stack;

    // Pointer used for traversing to new nodes in the tree
    private TreeNode current;

    // Keeps track of the last node we processed (returned to client)
    private TreeNode lastVisited;

    /**
     * Initialize iterator with root of BST
     * Time: O(h) - initial traversal to leftmost leaf
     * Space: O(h) - stack stores path from root to current
     */
    public PostorderIterator(TreeNode root) {
        stack = new Stack<>();
        current = root;
        lastVisited = null;
        // Start by finding the first node in postorder (leftmost leaf)
        findNextNode();
    }

    /**
     * Helper method to traverse as far left as possible, then right if no left exists
     * This helps us maintain the postorder traversal: Left -> Right -> Root
     * Time: O(h) where h is height of subtree
     */
    private void findNextNode() {
        while (current != null) {
            stack.push(current);
            current = current.left != null ? current.left : current.right;
        }
    }

    /**
     * Check if there are more nodes to process
     * Time: O(1)
     * @return true if there are more nodes to traverse
     */
    public boolean hasNext() {
        return !stack.isEmpty();
    }

    /**
     * Get the next node value in postorder sequence
     * Time: O(1) amortized
     * @return value of next node in postorder
     * @throws NoSuchElementException if no more nodes
     */
    public int next() {
        if (!hasNext()) {
            throw new java.util.NoSuchElementException();
        }

        while (true) {
            TreeNode node = stack.peek();
            // If we haven't visited right subtree yet and it exists
            if (node.right != null && lastVisited != node.right) {
                current = node.right;
                findNextNode();
                continue;
            }
            // Process current node
            lastVisited = stack.pop();
            return lastVisited.val;
        }
    }
}

InOrder

import java.util.Stack;

/**
 * Iterator implementation for Binary Search Tree that traverses in inorder (Left -> Root -> Right)
 * Provides a way to iterate through BST nodes without recursion, using O(h) space
 * where h is the height of the tree.
 */
public class InorderIterator {
    // Stack to store nodes in the path from root to current
    private Stack<TreeNode> stack;
    // Current node being processed
    private TreeNode current;

    /**
     * Initialize iterator with root of BST
     * Time: O(1) - just initializes data structures
     * Space: O(h) - stack will store at most h nodes where h is tree height
     */
    public InorderIterator(TreeNode root) {
        stack = new Stack<>();
        current = root;
        // No need to push root initially, will be handled in next()
    }

    /**
     * Check if there are more nodes to process
     * Time: O(1)
     * @return true if there are more nodes to traverse
     */
    public boolean hasNext() {
        return !stack.isEmpty() || current != null;
    }

    /**
     * Get the next node value in inorder sequence
     * Time: O(h) worst case when traversing to leftmost node, O(1) amortized
     * @return value of next node in inorder
     * @throws NoSuchElementException if no more nodes
     */
    public int next() {
        if (!hasNext()) {
            throw new java.util.NoSuchElementException();
        }

        // Traverse to leftmost node
        while (current != null) {
            stack.push(current);
            current = current.left;
        }

        TreeNode node = stack.pop();
        current = node.right;

        return node.val;
    }
}

PreOrder

import java.util.Stack;

/**
 * Iterator implementation for Binary Search Tree that traverses in preorder (Root -> Left -> Right)
 * Provides a way to iterate through BST nodes without recursion, using O(h) space
 * where h is the height of the tree.
 */
public class PreorderIterator {
    // Stack to store nodes that need to be processed
    private Stack<TreeNode> stack;

    /**
     * Initialize iterator with root of BST
     * Time: O(1) - just pushes root if it exists
     * Space: O(h) - stack will store at most h nodes where h is tree height
     */
    public PreorderIterator(TreeNode root) {
        stack = new Stack<>();
        if (root != null) {
            stack.push(root);
        }
    }

    /**
     * Check if there are more nodes to process
     * Time: O(1)
     * @return true if there are more nodes to traverse
     */
    public boolean hasNext() {
        return !stack.isEmpty();
    }

    /**
     * Get the next node value in preorder sequence
     * Time: O(1) - constant time operations
     * @return value of next node in preorder
     * @throws NoSuchElementException if no more nodes
     */
    public int next() {
        if (!hasNext()) {
            throw new java.util.NoSuchElementException();
        }

        TreeNode node = stack.pop();
        // Push right first so that left is processed first (LIFO)
        if (node.right != null) {
            stack.push(node.right);
        }
        if (node.left != null) {
            stack.push(node.left);
        }

        return node.val;
    }
}

Test Case

import org.junit.Test;
import static org.junit.Assert.*;
import java.util.ArrayList;
import java.util.List;

public class BSTIteratorTest {

    /**
     * Helper method to create a sample BST:
     *       4
     *      / \
     *     2   6
     *    / \ / \
     *   1  3 5  7
     */
    private TreeNode createSampleBST() {
        TreeNode root = new TreeNode(4);
        root.left = new TreeNode(2);
        root.right = new TreeNode(6);
        root.left.left = new TreeNode(1);
        root.left.right = new TreeNode(3);
        root.right.left = new TreeNode(5);
        root.right.right = new TreeNode(7);
        return root;
    }

    /**
     * Helper method to create an unbalanced BST:
     *   1
     *    \
     *     2
     *      \
     *       3
     */
    private TreeNode createUnbalancedBST() {
        TreeNode root = new TreeNode(1);
        root.right = new TreeNode(2);
        root.right.right = new TreeNode(3);
        return root;
    }

    @Test
    public void testInorderIterator() {
        TreeNode root = createSampleBST();
        InorderIterator iterator = new InorderIterator(root);
        List<Integer> result = new ArrayList<>();

        while (iterator.hasNext()) {
            result.add(iterator.next());
        }

        // Expected: [1, 2, 3, 4, 5, 6, 7]
        Integer[] expected = {1, 2, 3, 4, 5, 6, 7};
        assertArrayEquals(expected, result.toArray(new Integer[0]));
    }

    @Test
    public void testPreorderIterator() {
        TreeNode root = createSampleBST();
        PreorderIterator iterator = new PreorderIterator(root);
        List<Integer> result = new ArrayList<>();

        while (iterator.hasNext()) {
            result.add(iterator.next());
        }

        // Expected: [4, 2, 1, 3, 6, 5, 7]
        Integer[] expected = {4, 2, 1, 3, 6, 5, 7};
        assertArrayEquals(expected, result.toArray(new Integer[0]));
    }

    @Test
    public void testPostorderIterator() {
        TreeNode root = createSampleBST();
        PostorderIterator iterator = new PostorderIterator(root);
        List<Integer> result = new ArrayList<>();

        while (iterator.hasNext()) {
            result.add(iterator.next());
        }

        // Expected: [1, 3, 2, 5, 7, 6, 4]
        Integer[] expected = {1, 3, 2, 5, 7, 6, 4};
        assertArrayEquals(expected, result.toArray(new Integer[0]));
    }

    @Test
    public void testEmptyTree() {
        TreeNode emptyRoot = null;

        InorderIterator inorder = new InorderIterator(emptyRoot);
        assertFalse("Empty tree should have no next element for inorder", inorder.hasNext());

        PreorderIterator preorder = new PreorderIterator(emptyRoot);
        assertFalse("Empty tree should have no next element for preorder", preorder.hasNext());

        PostorderIterator postorder = new PostorderIterator(emptyRoot);
        assertFalse("Empty tree should have no next element for postorder", postorder.hasNext());
    }

    @Test
    public void testUnbalancedTree() {
        TreeNode root = createUnbalancedBST();

        // Test inorder
        InorderIterator inorder = new InorderIterator(root);
        List<Integer> inorderResult = new ArrayList<>();
        while (inorder.hasNext()) {
            inorderResult.add(inorder.next());
        }
        assertArrayEquals(new Integer[]{1, 2, 3}, inorderResult.toArray(new Integer[0]));

        // Test preorder
        PreorderIterator preorder = new PreorderIterator(root);
        List<Integer> preorderResult = new ArrayList<>();
        while (preorder.hasNext()) {
            preorderResult.add(preorder.next());
        }
        assertArrayEquals(new Integer[]{1, 2, 3}, preorderResult.toArray(new Integer[0]));

        // Test postorder
        PostorderIterator postorder = new PostorderIterator(root);
        List<Integer> postorderResult = new ArrayList<>();
        while (postorder.hasNext()) {
            postorderResult.add(postorder.next());
        }
        assertArrayEquals(new Integer[]{3, 2, 1}, postorderResult.toArray(new Integer[0]));
    }

    @Test(expected = java.util.NoSuchElementException.class)
    public void testInorderNoSuchElement() {
        TreeNode root = new TreeNode(1);
        InorderIterator iterator = new InorderIterator(root);
        iterator.next(); // Get the only element
        iterator.next(); // Should throw NoSuchElementException
    }

    @Test(expected = java.util.NoSuchElementException.class)
    public void testPreorderNoSuchElement() {
        TreeNode root = new TreeNode(1);
        PreorderIterator iterator = new PreorderIterator(root);
        iterator.next(); // Get the only element
        iterator.next(); // Should throw NoSuchElementException
    }

    @Test(expected = java.util.NoSuchElementException.class)
    public void testPostorderNoSuchElement() {
        TreeNode root = new TreeNode(1);
        PostorderIterator iterator = new PostorderIterator(root);
        iterator.next(); // Get the only element
        iterator.next(); // Should throw NoSuchElementException
    }
}