Ac solution code


  • 0

    Solution1. Flatten Tree to LinkedList in inorder sequence - Average O(1) time; O(1) space

    As the question description requires to output the 'number', not the node, so we can modify the tree structure.
    The basic idea is to flatten Tree to LinkedList in-place in inorder sequence: use TreeNode's right as LinkedNode's next.

    I implement two solutions for flatten function:

    1-1 Morris inorder Traversal: Runtime = O(2n)=O(n); Space = O(1)

    Morris traversal takes use of root's left's right-most leaf to save the root(leaf's next), so that it knows the next step for the whole traversal procedure. Regular Morris algorithm reverts leaf's right to null.

    How to set the node's next?

    last's next(right):

    1. if last is not a leaf: last's next = the left-most kid of last's right kid

    2. if last is a leaf: last's next = last's parent

    Time complexity: O(2n) = O(n)

    Each edge of the tree is visited twice, so the time complexity = O(2n) = O(n)

    1-2 Stack: Runtime = O(n); Space = O(n)

    JAVA Code:

    public class BSTIterator {
        private TreeNode curNode = null;
        public BSTIterator(TreeNode root) {
        	flatten(root);
        }
    
        /** @return whether we have a next smallest number */
        public boolean hasNext() {
        	return curNode != null;
        }
     
        /** @return the next smallest number */
        public int next() {
            if (!hasNext()) return Integer.MAX_VALUE;
            int val = curNode.val;
            curNode = curNode.right;
            return val;
        }
        
        // Solution1-1. Morris inorder Traversal: Time = O(2n)=O(n); Space = O(1)
        // Flatten the tree to linkedList in-place with inorder sequence, so that it will be in the ascending order    
        public void flatten(TreeNode root) {
            TreeNode temp = null, last = null;
            
            while(root != null) {
                if(root.left != null) {
                    temp = root.left;
                    while(temp.right != null && temp.right != root)// Find right-most leaf of root's left kid
                        temp = temp.right;
                    
                    if(temp.right != null) {
                        last = root;
                        root = root.right;// Switch to root's right part
                    } else {
                        temp.right = root;// right-most leaf of root's left kid, set its right to root
                        root = root.left;// Switch to root's left part
                    }
                }else {// left leaf: root's left is null
                    if (curNode == null)// Set the left-most node as the head of the linkedlist
                    	curNode = root;
                    
                    // last's next(right):
                    // 1) if last is not a leaf: last's next = the left-most kid of last's right kid
                    // 2) if last is     a leaf: last's next = last's parent                
                    if (last != null) 
                    	last.right = root;            
                    last = root;
                    root = root.right; // Switch to root's right part
                }
            }
        }
        
        //Solution1-2. Stack: Time = O(n); Space = O(n)
        // Flatten the tree to linkedList in-place with inorder sequence, so that it will be in the ascending order
        void flatten_stack(TreeNode root) {
            Stack<TreeNode> stack = new Stack<TreeNode>();
            TreeNode cur = root;
            TreeNode last = null;
            while (cur != null || !stack.isEmpty()) {
                while (cur != null) {
                	stack.push(cur); 
                	cur = cur.left;
                }
                cur = stack.pop();
                if (curNode == null)// Set the left-most node as the head of the linkedlist
                	curNode = cur;
                
                // last's next(right):
                // 1) if last is not a leaf: last's next = the left-most kid of last's right kid
                // 2) if last is     a leaf: last's next = last's parent
                if (last != null) 
                	last.right = cur;            
                last = cur;
                cur = cur.right;
            }
        }
    }
    

    Solution2. Use Stack without changing tree structure - Special thanks to xcv58

    Time: hasNext=O(1), next=O(h);

    Space: O(h): O(n) space totally; O(h) increased space for one time: because leftBranch must be shorter than h.

    public class BSTIterator {
    	private Stack<TreeNode>stack = null;
        public BSTIterator(TreeNode root) {
            stack = new Stack<TreeNode>();
            pushLeftBranch(root);// Push the left branch of the node
        }
    
        /** @return whether we have a next smallest number */
        public boolean hasNext() {
            return !stack.isEmpty();
        }
    
        /** @return the next smallest number */
        public int next() {
        	if (!hasNext()) return Integer.MAX_VALUE;
        	TreeNode node = stack.pop();
        	pushLeftBranch(node.right);// Push the left branch of the node's right kid
        	return node.val;
        }
        
        private void pushLeftBranch(TreeNode node) {
        	while (node != null) {
        		stack.push(node);
        		node = node.left;
        	}
        }
    }
    

Log in to reply
 

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