[Java] a solution of 15 lines


  • 9
    L

    The idea for the solution is composed of the following points.

    1. The solution uses a stack to keep track at most the next "h" (height of tree)
      elements for "next()" calls.
    2. The top of the stack is the current
      minimum element.
    3. At every "next()" call, we need to refresh the
      stack by populating the stack with all the left nodes up to the
      leaf, starting from the right node of the current minimum node.

    The complexity of the "hasNext()" is O(1). While the "next()" needs to refresh the stack, which in the best case (leaf) takes constant time, and in the worst case, it would take up to "h" steps. The overall cost of "next()" is then amortized over the number of nodes. I don't have the precise proof, but it seems to be O(1) on average.

    Here is the code. One trick is that I extract the refreshing into a function which can be used in the constructor as well, so that the code is more concise.

    Stack<TreeNode> stack = new Stack<TreeNode>();
    
    private void refreshStack(TreeNode iter){
         while(iter != null){
         	stack.push(iter);
         	iter = iter.left;
         }
    }
    	
    public BSTIterator(TreeNode root) {
       this.refreshStack(root);
    }
    
    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return !(stack.isEmpty());
    }
    
    /** @return the next smallest number */
    public int next() {
        TreeNode node = stack.pop();
        if(node != null){
        	this.refreshStack(node.right);
        	return node.val;
        }
        
        return -1; // should throw exception here.
    }

  • 0
    G

    I will throw an exception rather than return a magic value.


  • 0
    L

    Man you did a great job, here's a advice that makes your code more clean.

    In the next() function, you do not need to judge if node is null, because the nodes in the stack have no chance to be null.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.