class BSTIterator {
private:
stack<TreeNode*> st;
void update(TreeNode* cur) {
for (; cur; cur = cur>left) {
st.push(cur); // push candidate left nodes into stack all the way
}
}
public:
BSTIterator(TreeNode *root) {
update(root); // add root to stack
}
/** @return whether we have a next smallest number */
bool hasNext() {
return st.size(); // has next as long as stack is not empty
}
/** @return the next smallest number */
int next() {
TreeNode* cur = st.top(); st.pop();
update(cur>right); // get leftmost availabe node and add in its candidate right node
return cur>val;
}
};
Clean C++ solution


The problem has the following requirement:
"Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree."
It seems that functionupdate(TreeNode * node)
has O(h) time complexity (i.e., traverse to leftmost leaf of input node), so I thinknext()
also has O(h) time complexity instead of O(1) since it callsupdate(TreeNode * node)
to find the leftmost leaf of its right subtree.Am I missing something?

@zzg_zzm
Suppose we have N as the total number of tree node.By calling
update()
innext()
, each tree node will get pushed into the stack exactly once, and get popped out from the stack exactly once.
Also,next()
will get called N times in total. So the average run time is still O(1).