A C++ solution with lg(n) space but doubt the time is O(1)


  • 2
    G

    The basic idea is to manually maintain a stack for the in-order traversal, which is accepted. But I am not sure why the time is O(1)

     class BSTIterator {
    private:
    	stack<TreeNode*> track;
    
    public:
    	BSTIterator(TreeNode *root)
    	{
    		if (root)
    		{
    			TreeNode *p = root;
    			while (p)
    			{
    				track.push(p);
    				p = p->left;
    			}
    		}
    	}
    
    	/** @return whether we have a next smallest number */
    	bool hasNext() 
    	{
    		return track.size() > 0;
    	}
    
    	/** @return the next smallest number */
    	int next() 
    	{
    		if (track.size() == 0) throw exception();
    		TreeNode* pTop = track.top();
    		track.pop();
    		
    		TreeNode* p = pTop->right;
    		while (p)
    		{
    			track.push(p);
    			p = p->left;
    		}
    
    		return pTop->val;
    	}
    };

  • 7
    S

    It is 'amortized O(1)' or 'O(1) on average', because when we call next() on all elements in a BST, each node is pushed into and pop from the stack exactly once.


  • 0
    I

    Same idea, but more concise

    class BSTIterator {
    public:
        stack<TreeNode*> s;
        
        BSTIterator(TreeNode *root) {
            update(root);
        }
        
        void update(TreeNode* root){
            while(root){
                s.push(root);
                root=root->left;
            }
        }
    
        /** @return whether we have a next smallest number */
        bool hasNext() {
            return !s.empty();
        }
    
        /** @return the next smallest number */
        int next() {
            TreeNode* t = s.top(); s.pop();
            update(t->right);
            return t->val;
        }
    };

Log in to reply
 

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