Simple C++ solution based on stack WITH explanation


  • 0
    U

    Basic idea: use a stack to keep trace of the ancestors of current node.
    The next smallest node could be found as the following rules:

    • [0] start by keep go to left son from root: root->left->...->left
    • [1] if the current node has a right son, go to current->right->left->...->left
    • [2] else pop the stack until the father node has a bigger value (if cannot find such node, we've already finished)

    For example, the tree [5, 2, 6, 1, 4, null, null, null, null, 3] would be visited by applying rules like [0]->[2]->[1]->[2]->[2]->[1].

    class BSTIterator {
    private:
        stack<TreeNode*> st;
        TreeNode* node;
        bool flag;
    public:
        BSTIterator(TreeNode *root) {
            while (root && root->left) {
                st.push(root);
                root = root->left;
            }
            node = root;
            flag = root;
        }
    
        /** @return whether we have a next smallest number */
        bool hasNext() {
            return flag;
        }
    
        /** @return the next smallest number */
        int next() {
            int tmp = node->val;
            if (node->right) {
                st.push(node);
                node = node->right;
                while (node->left) {
                    st.push(node);
                    node = node->left;
                }
                return tmp;
            }
            while (st.size() && st.top()->val < node->val) {
                node = st.top();
                st.pop();
            }
            if (!st.size()) {
                flag = false;
                return tmp;
            }
            node = st.top();
            st.pop();
            return tmp;
        }
    };
    

  • 0

    The function next() seems to have O(h) time complexity instead of O(1). Am I missing something?


  • 0
    U

    @zzg_zzm The Problem said:
    "The next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree."
    We have n nodes and each of them goes into/out of the stack at most once, and the stack size is at most h. Thus the average running time of next() is O(1) and uses O(h) memory.


  • 0

    Thanks. I see. It is the AVERAGE time complexity required by this problem, as I was thinking about the worse cases (when a node has a huge right subtree which takes O(log (subtree height)) time).

    I initially also thought about simply using In-order traversal, but I was stuck at the time complexity requirement for next().


Log in to reply
 

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