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

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