Problem

Implement the BSTIterator class that represents an iterator over the in-order traversal of a binary search tree (BST):

  • BSTIterator(TreeNode root) Initializes an object of the BSTIterator class. The root of the BST is given as part of the constructor. The pointer should be initialized to a non-existent number smaller than any element in the BST.
  • boolean hasNext() Returns true if there exists a number in the traversal to the right of the pointer, otherwise returns false.
  • int next() Moves the pointer to the right, then returns the number at the pointer.

Notice that by initializing the pointer to a non-existent smallest number, the first call to next() will return the smallest element in the BST.

You may assume that next() calls will always be valid. That is, there will be at least a next number in the in-order traversal when next() is called.

Your Solution will be called like this:

Solution i = new Solution(root);
while (i.hasNext()) System.out.print(i.next());

Examples

Example 1:

     7
   /   \
  3     15
       /   \
      9    20

Input
["BSTIterator", "next", "next", "hasNext", "next", "hasNext", "next", "hasNext", "next", "hasNext"]
[ [7, 3, 15, null, null, 9, 20]], [], [], [], [], [], [], [], [], [] ]
Output
[null, 3, 7, true, 9, true, 15, true, 20, false]

Explanation
BSTIterator bSTIterator = new BSTIterator([7, 3, 15, null, null, 9, 20]);
bSTIterator.next();    // return 3
bSTIterator.next();    // return 7
bSTIterator.hasNext(); // return True
bSTIterator.next();    // return 9
bSTIterator.hasNext(); // return True
bSTIterator.next();    // return 15
bSTIterator.hasNext(); // return True
bSTIterator.next();    // return 20
bSTIterator.hasNext(); // return False

Follow up

  • Could you implement next() and hasNext() to run in average O(1) time and use O(h) memory, where h is the height of the tree?

Solution

The code is similar to iterative inorder traversal - Binary Tree Inorder Traversal.

We should also understand what is Iterator Design Pattern.

Method 1 - Naive Solution

A straightforward approach to a binary search tree iterator would involve creating a List (or similar linear data structure) during the construction of the BSTIterator. This method requires traversing the entire BST and adding each node value to the list. However, this approach has two significant drawbacks:

  • It requires the BSTIterator’s constructor to traverse the whole tree upfront,
  • and it consumes more memory by storing every node (in the worst case) in a separate data structure.

In Big O terms, where n is the number of nodes in the tree:

  • Constructor Run Time: O(n)
  • next() Run Time: O(1)
  • hasNext() Run Time: O(1)
  • Extra Space: O(n)

This approach does not address the follow-up requirement to only use O(h) memory, where h is the height of the tree.

Furthermore, a recursive solution may not be ideal, as it requires additional node traversal when adding values to the recursive stack. For example, at the root, you’d need to traverse to the leftmost element during recursion. Therefore, an iterative solution is simpler and more efficient.

Method 2 - Iterative Inorder Traversal Using Stack

An improved approach is to track only the path to the next smallest node. For example, for the given BST, when initializing your BSTIterator, instead of traversing the entire tree, you would go directly to the next smallest node (in this example, 1), pushing nodes onto a stack as you proceed.

Code

Java
Without code reuse
public class BSTIterator {
  Stack <TreeNode> stack;

  public BSTIterator(TreeNode root) {
    stack = new Stack <TreeNode> ();
    while (root != null) {
      stack.push(root);
      root = root.left;
    }
  }

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

  public int next() {
    TreeNode node = stack.pop();
    int result = node.val;
    if (node.right != null) {
      node = node.right;
      while (node != null) {
        stack.push(node);
        node = node.left;
      }
    }
    return result;
  }
}
With code reuse

We can reuse some part of filling the left nodes in stack using fillStack function.

class BSTIterator {
    private Stack<TreeNode> stack;
    public BSTIterator(TreeNode root) {
        stack = new Stack<>();
        fillStack(root);
    }

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

    /**
     * @return the next smallest number
     */
    public int next() {
	    if(stack.isEmpty()) {
		    return -1;
	    }
		TreeNode top = stack.pop();
		fillStack(top.right);
		return top.val;        
    }

    /**
     * Fill stack with min values
     *
     * @param node curr node to begin with
     */
    private void fillStack(TreeNode curr) {
        while (curr != null) {
            stack.push(curr);
            curr = curr.left;
        }
    }
}

Complexity

In Big O terms, the costs for our improved solution, where h is the height of the tree, are:

  • Constructor Run Time: O(h)
  • next() Amortized Run Time: O(1)
  • hasNext() Run Time: O(1)
  • Extra Space: O(h)

This solution distributes the computational work more evenly. It provides the same amortized runtime for next(), shortens the constructor’s runtime, and requires less additional space.