LeetCode – Binary Search Tree Iterator (Java)

by

Problem

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST. Calling next() will return the next smallest number in the BST. Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.

Java Solution

The key to solve this problem is understanding the features of BST. Here is an example BST.

binary-search-tree

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
 
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;
	}
}

15 thoughts on “LeetCode – Binary Search Tree Iterator (Java)”

  1. storing nodes in a queue will incur O(n) space complexity where n is the number of nodes. Here Question ask to implement in O(h) space complexity ,h is the height of the tree.In this solution number of nodes in the stack cannot exceed h anytime.

  2. I am new in tree, but from my understanding In order traversal of a BST is sorted, if we just do in order traversal and put them in queue, and them just use that queue for popping the value.. what will be the difference and the code will be simpler, because at the end you are also doing in order traversal

  4. actually the way iterator works (fail-fast) is to throw a ConcurrentModificationException if any insertion/deletion occurs during iterating. (i.e. once the iterator() method is called and an iterator is returned, trying to use that iterator after modifying the collection will result in ConcurrentModificationException)
    Even if you choose the to make the iterator fail-safe, you will copy the collection and ignore any changes to it during iteration.
    Keep track of elements in both tree and stack is cumbersome, another solution would be better (i.e. make every node know its parent and then you can go back and forth)

  6. Agreed, it should check if stack is empty and throw a NoSuchElementException per the iterator interface standards.

  7. class BSTIterator{
    Stack stack;
    public BSTIterator(Node n){
    stack=new Stack();
    add(n);
    }
    private void add(Node n){
    if(n.left!=null){
    stack.push(n.left);
    Node curr=n.left.right;
    while(curr!=null){ stack.push(curr); curr=curr.right;}
    }
    }
    public Node next(){
    if(stack.size()>0){
    Node top=stack.pop();
    add(top);
    return top;
    }
    return null;
    }
    public boolean hasNext(){
    if(stack.size()>0) return true;
    return false;
    }
    }

  10. Does this work if the iterator is initialized at the root?

    Does the iterator need to be at the left most node?

  12. I think the key phrase is ‘on average’. If you you iterate over all tree, you visit each node only once. Hence, overall time will be O(n).

  13. if (node.right != null) {
    node = node.right;
    while (node != null) {
    stack.push(node);
    node = node.left;
    }
    }

    This can be simplified to be:

    TreeNode child = node.right;
    while (child != null) {
    stack.push(child);
    child = child.left;
    }

  15. Good solution, but this would not handle modifications to the BST. You would need to update the stack to keep track of any inserts or delete’s that could happen.

Leave a Comment