You are given a task to finish the class and implement two functions.

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, while calling hasNext() will return whether we have a next smallest number.

bst-tree

BSTIterator iterator = new BSTIterator(root);
iterator.next();    // return 3
iterator.next();    // return 7
iterator.hasNext(); // return true
iterator.next();    // return 9
iterator.hasNext(); // return true
iterator.next();    // return 15
iterator.hasNext(); // return true
iterator.next();    // return 20
iterator.hasNext(); // return false

Here's the solution:

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class BSTIterator {
    
    private Stack <TreeNode> stack = new Stack<>();

    public BSTIterator(TreeNode root) {
        while(root != null){
            stack.push(root);
            root = root.left;  
        }
        
    }
    
    /** @return the next smallest number */
    public int next() {
        TreeNode curr = stack.peek();
        TreeNode node = curr;
        if(node.right == null){
            node = stack.pop();
            while(!stack.isEmpty() && stack.peek().right == node){
                node = stack.pop();
            }
            
        } else{
            node = node.right;
            while(node != null){
                stack.push(node);
                node = node.left;
            }
        }
        
        return curr.val;
    }
    
    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return !stack.isEmpty();
    }
}

/**
 * Your BSTIterator object will be instantiated and called as such:
 * BSTIterator obj = new BSTIterator(root);
 * int param_1 = obj.next();
 * boolean param_2 = obj.hasNext();
 */

We are using stack to implement these two functions. Since the root is a BST tree, the smallest number will be the leftmost node. Hence, when we initialize BSTIterator class, we keep pushing the node into the stack until the node is null.

In hasNext() function, we first retrieve the "smallest" node by calling stack.peek(). Once we call next() function, we need to readjust the number order in stack. Here, we need to evaluate these two conditions:

  1. If the current node has right node, then the next smallest node will be the leftmost child node of the right node.
  2. If the current node has no right node, then the next smallest node will be the first left-turning node of current path.

Since visiting all the nodes take O(n) time, hence visiting each node will require O(1) time.


That's it! See you all next time!

Post was published on , last updated on .

Like the content? Support the author by paypal.me!