O(1) average time O(1) space with Morris traverse


  • 0
    C

    The idea is to store next TreeNode of Morris traverse, then resume from there when next() is called
    Morris traverse is O(n) for a tree with n nodes, when next() is called n times, only one complete Morris traverse is needed (break up by n times call to morris() method though), hence average time is O(1) for next() method

    public class BSTIterator {
        private TreeNode next;
        private int[] buf;
        private int p;
        private int size;
        private int visited;
    
        private int count(TreeNode root) {
            return root == null ? 0 : 1 + count(root.left) + count(root.right);   
        }
        
        private void morris(TreeNode r) {
            if(r != null) {
                int id = 0;
                while(r != null) {
                    if(r.left == null) {
                        if(id < buf.length) {
                            buf[id++] = r.val;
                        }else {
                            next = r;
                            break;
                        }
                        r = r.right;
                    }else {
                        TreeNode pre = r.left;
                        while(pre.right != null && pre.right != r)
                            pre = pre.right;
                        if(pre.right == null) {
                            pre.right = r;
                            r = r.left;
                        }else {
                            if(id < buf.length) {
                                buf[id++] = r.val;
                            }else {
                                next = r;
                                break;
                            }
                            pre.right = null;
                            r = r.right;
                        }
                    }
                }
                p = 0;
            }
        }
        
        public BSTIterator(TreeNode root) {
            if(root != null) {
                size = count(root);
                buf = new int[1];
                morris(root);
                p = 0;
                visited = 0;
            }
        }
    
        /** @return whether we have a next smallest number */
        public boolean hasNext() {
            return visited < size;
        }
    
        /** @return the next smallest number */
        public int next() {
            if(hasNext()) {
                if(p >= buf.length)
                    morris(next);
                visited++;
                return buf[p++];
            }
            return -1;
        }
    }
    

Log in to reply
 

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