LeetCode – Binary Search Tree Iterator (Java)

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;
	}
}
Category >> Algorithms >> Interview >> Java  
If you want someone to read your code, please put the code inside <pre><code> and </code></pre> tags. For example:
<pre><code> 
String foo = "bar";
</code></pre>

  1. Abhi on 2015-6-17

    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.

  2. antonio081014 on 2015-8-25

    Well, I don’t think your next() takes O(1), I think it takes O(h) instead.

  3. Lookuptable on 2016-1-19

    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;
    }

  4. Defiler on 2016-1-24

    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).

  5. antonio081014 on 2016-1-24

    Right. Kindly miss that.

  6. Aditya Vutukuri on 2016-6-7

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

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

  7. Subodh Karwa on 2016-11-10

    next() should check if stack is empty before popping to avoid NullPointerException

  8. jaqen h'ghar on 2016-12-3

    Usually next() should be used only after checking hasNext().

  9. Iram22 on 2017-2-13

    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;
    }
    }

Leave a comment

*