173. Binary Search Tree Iterator

For this problem, if don't consider the follow up limitation, the solution is very easy:

The time complexity is O(1), the space complexity is O(n), n is the total node number of the tree.

class BSTIterator {
    Queue<Integer> queue = new LinkedList<>();
    public BSTIterator(TreeNode root) {
        inOrder(root);
    }
    
    private void inOrder(TreeNode root){
        if(root==null)
            return;
        inOrder(root.left);
        queue.offer(root.val);
        inOrder(root.right);
    }
    
    public int next() {
        return queue.poll();
    }
    
    public boolean hasNext() {
        return !queue.isEmpty();
    }
}

 

But the problem have a follow up:

Could you implement next() and hasNext() to run in average O(1) time and use O(h) memory, where h is the height of the tree?

I can think of a solution with next() and hasNext() to run in average O(1) time, and use O(1) memory, but this solution need to change the tree's structure. This solution conver the tree to a double direction linked list, inspired by: 426. Convert Binary Search Tree to Sorted Doubly Linked List

class BSTIterator {

    private TreeNode head,  pre = new TreeNode();
    public BSTIterator(TreeNode root) {
        head = root;
        if(head!=null){
            while(head.left!=null)
                 head = head.left;   //1. find the smallest element
        inOrder(root);
        pre.right = head;
        head.left = pre;
      }
    }
        
    private void inOrder(TreeNode root){ 
        if(root == null)
            return;
        inOrder(root.left);
        pre.right = root;
        root.left = pre;
        pre = root;
        inOrder(root.right);
    }
    int count =0;
    public int next() {
        count++;
        pre=pre.right;
        return pre.val;
    }
    
    public boolean hasNext() {
        return count>0 && pre.right!=head||count==0;
    }
}

The following is other's solution, the space complexity is O(h), but either next() or hasNext() cannot have O(1) time complexity:

class BSTIterator {
    private Stack<TreeNode> stk = new Stack<>();

    public BSTIterator(TreeNode root) {
        pushLeft(root);
    }

    private pushLeft(TreeNode root) {
        while (root != null) {
            stk.push(root);
            root = root.left;
        }
    }

    public int next() {
        TreeNode temp = stk.pop();
        pushLeft(temp.right)
        return temp.val;
    }

    public boolean hasNext() {
        return !stk.isEmpty();
    }
}