Ideal Solution using Stack (Java)


  • 121
    S

    My idea comes from this: My first thought was to use inorder traversal to put every node into an array, and then make an index pointer for the next() and hasNext(). That meets the O(1) run time but not the O(h) memory. O(h) is really much more less than O(n) when the tree is huge.

    This means I cannot use a lot of memory, which suggests that I need to make use of the tree structure itself. And also, one thing to notice is the "average O(1) run time". It's weird to say average O(1), because there's nothing below O(1) in run time, which suggests in most cases, I solve it in O(1), while in some cases, I need to solve it in O(n) or O(h). These two limitations are big hints.

    Before I come up with this solution, I really draw a lot binary trees and try inorder traversal on them. We all know that, once you get to a TreeNode, in order to get the smallest, you need to go all the way down its left branch. So our first step is to point to pointer to the left most TreeNode. The problem is how to do back trace. Since the TreeNode doesn't have father pointer, we cannot get a TreeNode's father node in O(1) without store it beforehand. Back to the first step, when we are traversal to the left most TreeNode, we store each TreeNode we met ( They are all father nodes for back trace).

    After that, I try an example, for next(), I directly return where the pointer pointing at, which should be the left most TreeNode I previously found. What to do next? After returning the smallest TreeNode, I need to point the pointer to the next smallest TreeNode. When the current TreeNode has a right branch (It cannot have left branch, remember we traversal to the left most), we need to jump to its right child first and then traversal to its right child's left most TreeNode. When the current TreeNode doesn't have a right branch, it means there cannot be a node with value smaller than itself father node, point the pointer at its father node.

    The overall thinking leads to the structure Stack, which fits my requirement so well.

    /**
     * Definition for binary tree
     * public class TreeNode {
     *     int val;
     *     TreeNode left;
     *     TreeNode right;
     *     TreeNode(int x) { val = x; }
     * }
     */
    
    public class BSTIterator {
        
        private Stack<TreeNode> stack;
        public BSTIterator(TreeNode root) {
            stack = new Stack<>();
            TreeNode cur = root;
            while(cur != null){
                stack.push(cur);
                if(cur.left != null)
                    cur = cur.left;
                else
                    break;
            }
        }
    
        /** @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();
            TreeNode cur = node;
            // traversal right branch
            if(cur.right != null){
                cur = cur.right;
                while(cur != null){
                    stack.push(cur);
                    if(cur.left != null)
                        cur = cur.left;
                    else
                        break;
                }
            }
            return node.val;
        }
    }
    
    /**
     * Your BSTIterator will be called like this:
     * BSTIterator i = new BSTIterator(root);
     * while (i.hasNext()) v[f()] = i.next();
     */

  • 0
    G

    hasNext() is O(1) for sure. How can we prove next() has average O(1) complexity?


  • 9
    S

    If the returned TreeNode doesn't have a right branch, it is O(1). If it has a right branch, it will traversal until its right child's left-most TreeNode. Now, considering the code under "// traversal right branch". After I use next() to go through the entire tree once, this part of code will traversal each node with a right child once. The total run time of this part is O(n). And since I go through the entire tree, the average next() run time is O(n) / n = O(1)


  • 0
    G

    good explanation! thank you!


  • 16
    A

    Similar idea using stack, but I make use of the answer from Binary Tree Inorder Traversal.

     public class BSTIterator {
            
            private TreeNode current; 
            private Stack<TreeNode> stack;
            public BSTIterator(TreeNode root) {
                current = root;
                stack = new Stack<TreeNode>();
            }
        
            /** @return whether we have a next smallest number */
            public boolean hasNext() {
                return !stack.isEmpty() || current!=null;
            }
        
            /** @return the next smallest number */
            public int next() {
                while (!stack.isEmpty() || current!=null) {
                    if (current!=null) {
                        stack.push(current);
                        current = current.left;
                    } else {
                        current = stack.peek().right;
                        break;
                    }
                }
                TreeNode node = stack.pop();
                return node.val;
            }
        }
    

  • 2
    S

    Thank you for sharing. I made a simplified version http://sprunge.us/SOEQ of your solution.


  • 0
    H

    nice solution!


  • 0
    Z

    good explanation!


  • 0
    F

    I like this solution


  • 0
    T

    I liked your solution and explanation about how the average time complexity is O(1). Good work.


  • 16
    C

    Nice explanation! Thank you. I rewrite it to be a little bit shorter.

    public class BSTIterator {
    
    Stack<TreeNode> stack;
    public BSTIterator(TreeNode root) {
        stack = new Stack<>();
        fillStack(root);
    }
    
    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return !stack.isEmpty();
    }
    
    /** @return the next smallest number */
    public int next() {
        TreeNode curNode = stack.pop();
        fillStack(curNode.right);
        return curNode.val;
    }
    
    private void fillStack(TreeNode root){
        while(root != null){
            stack.push(root);
            root = root.left;
        }
    }
    

    }


  • -1
    C

    @siyang2

    Nice thought!


  • 0
    K

    @siyang2

    yo man. thx for ur sharing first!
    For helping understanding why it only cause O(1) for next().
    Suppose we have a tree like: node A have B as right child, and B have 4 depth for the nodes who only have left childs (like a link list.)

    A
           B
       C
    D
    

    E

    we do the next() for 5 times, the time should be 1+4+1+1+1 = 9. if B's left child list size is m. total size would be m+1. Then the total time: 1+m+1+1+...+1 = 2m. which is O(m) for m operations. So average each time is O(1)


  • 0
    D

    One line improvement for your "// traversal right branch"& BSTIterator constructor:

    for (auto temp = n->right; temp; temp = temp->left) s.push(temp);
    

  • 0
    Stack<TreeNode> stack;
    public BSTIterator(TreeNode root) {
        stack = new Stack();
        TreeNode cur = root;
        while(cur != null){
            stack.push(cur);
            cur = cur.left;
        }
    }
    
    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return !stack.isEmpty();        
    }
    
    /** @return the next smallest number */
    public int next() {
        TreeNode tmp = stack.pop();
        TreeNode cur = tmp.right;
        while(cur != null){
            stack.push(cur);
            cur = cur.left;
        }
        return tmp.val;
    }

  • 0
    S

    I have a doubt!

    If the smallest node is the left most node then the 2nd most smallest element should be its parent right!?? why do we check for the right node? and the children on the left of the left of the right node!


  • 0
    B

    @sushant_gupta Because all the right child of current node are bigger than current node value but smaller than any other nodes, and thus should be push into the stack.


  • 0
    B

    I think in the next() method, it is better to check hasNext() first. Otherwise, may incur NullPointerException.


  • 0

    @ke8511369
    Once you need to push the entire left branch of the one node, it is O(h) operation.
    I think you may misunderstand the definition of the big O.
    O(1) means the complexity is constant while the size of the problem grows up. But for this problem, the cost of the operation mentioned above will grow with the height of the BST.


  • 0
    Z

    Nice explanation! Similar idea using Stack. Average time for geting next element is O(1) on average and the space complexity would be O(h) because we only store one element at every level.

    public class BSTIterator {
        
        Stack<TreeNode> stack;
    
        public BSTIterator(TreeNode root) {
            stack = new Stack<>();
            pushAll(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();
            pushAll(node.right);
            return node.val;
        }
        
        private void pushAll(TreeNode root){
            if(root == null) return;
            TreeNode curNode = root;
            while(curNode != null){
                stack.push(curNode);
                curNode = curNode.left;
            }
        }
    }
    

Log in to reply
 

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