[C++] Clean Code - O(n) time O(lg n) space - BinaryTree Iterator


  • 15

    O(n) time O(n) space - Inorder Vector

    class Solution {
    public:
        bool findTarget(TreeNode* root, int k) {
            vector<int> nums;
            inorder(root, nums);
            return findTargetInSortedArray(nums, k);
        }
    
    private:
        void inorder(TreeNode* node, vector<int>& nums) {
            if (!node) return;
            inorder(node->left, nums);
            nums.push_back(node->val);
            inorder(node->right, nums);
        }
    
        bool findTargetInSortedArray(vector<int> a, int target) {
            for (int i = 0, j = a.size() - 1; i < j;) {
                int sum = a[i] + a[j];
                if (sum == target) {
                    return true;
                }
                else if (sum < target) {
                    i++;
                }
                else {
                    j--;
                }
            }
    
            return false;
        }
    };
    

    O(n) time O(lg n) space - BinaryTree Iterator

    class BSTIterator {
        stack<TreeNode*> s;
        TreeNode* node;
        bool forward;
    public:
        BSTIterator(TreeNode *root, bool forward) : node(root), forward(forward) {};
        bool hasNext() {
            return node != nullptr || !s.empty();
        }
        int next() {
            while (node || !s.empty()) {
                if (node) {
                    s.push(node);
                    node = forward ? node->left : node->right;
                }
                else {
                    node = s.top();
                    s.pop();
                    int nextVal = node->val;
                    node = forward ? node->right : node->left;
                    return nextVal;
                }
            }
    
            return -1;  // never reach & not necessary
        }
    };
    class Solution {
    public:
        bool findTarget(TreeNode* root, int k) {
            if (!root) return false;
            BSTIterator l(root, true);
            BSTIterator r(root, false);
            for (int i = l.next(), j = r.next(); i < j;) {
                int sum = i + j;
                if (sum == k) {
                    return true;
                }
                else if (sum < k) {
                    i = l.next();
                }
                else {
                    j = r.next();
                }
            }
            return false;
        }
    };
    

    Simplified using operator overloading

    class BSTIterator {
        stack<TreeNode*> s;
        TreeNode* node;
        bool forward;
        int cur;
    public:
        BSTIterator(TreeNode *root, bool forward) : node(root), forward(forward) { (*this)++; };
        int operator*() { return cur; }
        void operator++(int) {
            while (node || !s.empty()) {
                if (node) {
                    s.push(node);
                    node = forward ? node->left : node->right;
                }
                else {
                    node = s.top();
                    s.pop();
                    cur = node->val;
                    node = forward ? node->right : node->left;
                    break;
                }
            }
        }
    };
    class Solution {
    public:
        bool findTarget(TreeNode* root, int k) {
            if (!root) return false;
            BSTIterator l(root, true);
            BSTIterator r(root, false);
            while (*l < *r) {
                if (*l + *r == k) return true;
                else if (*l + *r < k) l++;
                else r++;
            }
            return false;
        }
    };
    

  • 3
    Z

    Good solution!
    I think the second iterator approach is average O(lgn) space for a balanced tree. The worst case could be O(n). What if it is a linked list with only left child?


  • 0

    @zestypanda This is a good interview question, I guess that would be O(N)?


  • 7
    M

    @alexander
    This solution should be upvoted!!! O(N) time complexity and O(h) space complexity.
    The idea is pretty straightforward and I love the way the code was written.
    Note: This is actually a Snapchat onsite interview problem. The interviewer explicitly requested the space complexity should be better than O(N) as a follow up question.

    My java version inspire by alexander

    class Solution {
        class BSTIterator {
            private Stack<TreeNode> stack;
            private boolean isForward = true;
            public BSTIterator(TreeNode root, boolean flag) {
                this.isForward = flag;
                stack = new Stack<>();
                TreeNode node = root;
                while (node != null) {
                    stack.push(node);
                    node = isForward ? node.left : node.right;
                }
            }
    
            /** @return whether we have a next smallest number */
            public boolean hasNext() {
                return !stack.isEmpty();
            }
            public int peek() {
                return stack.peek().val;
            }
            /** @return the next smallest number */
            public int next() {
                TreeNode node = stack.pop();
                int val = node.val;
                node = isForward ? node.right : node.left;
                while (node != null) {
                    stack.push(node);
                    node = isForward ? node.left : node.right;
                }        
                return val;
            }
        }
        
        public boolean findTarget(TreeNode root, int k) {
            BSTIterator left  = new BSTIterator(root, true);
            BSTIterator right = new BSTIterator(root, false);
            while (left.hasNext() && right.hasNext()) {
                int l = left.peek(), r = right.peek();
                if (l >= r)  return false;
                if (l + r == k) return true;
                else if (l + r < k)
                    left.next();
                else
                    right.next();
            }
            return false;
        }    
    }
    

  • 0
    S

    @MichaelPhelps
    Thanks for sharing, but I noticed that return int may cause some issue, since the same int l and int r, may represent different node, if input is 1, 1, target is 2, your code will return false. Instead of int, we can return TreeNode, then we can compare two node, if they are the same meaning l == r, we can return false
    here is the code :

    class Solution {
        public boolean findTarget(TreeNode root, int k) {
            //next smallest and next largest, time o(n), space o(h)
            BSTIterator bstSmall = new BSTIterator(root, true);
            BSTIterator bstLarge = new BSTIterator(root, false);
            while (bstSmall.hasNext() && bstLarge.hasNext()) {
                TreeNode l = bstSmall.peek();
                TreeNode r = bstLarge.peek();
                
                if (l == r) {
                    return false;
                }
                if (l.val + r.val == k) {
                    return true;
                } 
                if (l.val + r.val > k) {
                    bstLarge.next();
                } else {
                    bstSmall.next();
                }
            }
            return false;
        }
        private class BSTIterator {
            private Stack<TreeNode> stack;
            private boolean flag;
            public BSTIterator(TreeNode root, boolean flag) {
                stack = new Stack<>();
                this.flag = flag;
                addToStack(root, flag);
            }
            
            private boolean hasNext() {
                return !stack.isEmpty();
            }
            private TreeNode next() {
                TreeNode cur = stack.pop();
                TreeNode nextNode = flag? cur.right : cur.left;
                addToStack(nextNode, flag);
                return cur;
            }
            private TreeNode peek() {
                return stack.peek();
            }
            private void addToStack(TreeNode root, boolean flag) {
                if (root == null) {
                    return;
                }
                while (root != null) {
                    stack.push(root);
                    root = flag? root.left : root.right;
                }
                return;
            }
        }    
    }
    

Log in to reply
 

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