Strict O(1) time for both hasNext() and next() : a destructive approach


  • 0
    C

    It is still feasible to achieve the time constraints, i.e. a strict O(1) as long as we can change the tree structure. The main idea is to flatten the tree nodes using its right child in order, pretty much like generating an ordered linked list conceptually.

    The following code (still can be improved) fails OJ as expected probably because the destructive change in its structure messes up with the program cleanup. However the output checks out (if you cout the result).

    /**
     * Definition for binary tree
     * struct TreeNode {
     *     int val;
     *     TreeNode *left;
     *     TreeNode *right;
     *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
     * };
     */
    class BSTIterator {
    public:
        BSTIterator(TreeNode *root) {
            // a destructive approach
            // node->right is modified to point to the node with the next larger number.
            // node->left is unchanged
            stack<TreeNode*> s;
            TreeNode *p = nullptr; // the current node already in sequence
            
            if (root)
                s.push(root);
                
            while (!s.empty())
            {
                TreeNode *q;
                while ((q = s.top()->left))
                    s.push(q);
                
                if (p) p->right = s.top();
    
                while (!((p = s.top())->right))
                {
                    s.pop();
                    if (s.empty()) // only for the last node on the stack
                        break;
                    else
                        p->right = s.top();
                }
                
                if (p->right)
                {
                    s.pop();
                    s.push(p->right);
                }
            }
            
            if ((it = root))
                while (it->left)
                    it = it->left;
        }
    
        /** @return whether we have a next smallest number */
        bool hasNext() {
            return it;
        }
    
        /** @return the next smallest number */
        int next() {
            int val = it->val;
            it = it->right;
            return val;
        }
    
    private:
        TreeNode* it;
    };
    
    /**
     * Your BSTIterator will be called like this:
     * BSTIterator i = BSTIterator(root);
     * while (i.hasNext()) cout << i.next();
     */

Log in to reply
 

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