LeetCode – Binary Search Tree Iterator (Java)

Category: Algorithms >> Interview >> Java   April 1, 2014

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

Related posts:

  1. Leetcode – Binary Tree Inorder Traversal (Java)
  2. LeetCode – Lowest Common Ancestor of a Binary Search Tree (Java)
  3. LeetCode – Validate Binary Search Tree (Java)
  4. LeetCode – Closest Binary Search Tree Value (Java)
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>
  • ajitesh kumar

    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.

  • shivam gupta

    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

  • thank you but i didn’t understand shit

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

  • Cat Racket
  • David McGuire

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

  • Iram22

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

  • jaqen h’ghar

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

  • Subodh Karwa

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

  • Aditya Vutukuri

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

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

  • Right. Kindly miss that.

  • Defiler

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

  • Lookuptable

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

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

  • Abhi

    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.