Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.
Calling next()
will return the next smallest number in the BST.
Note: next()
and hasNext()
should run in average O(1) time and uses O(h) memory, where h is the height of the tree.
Solution
public class BSTIterator {
private Stack<TreeNode> stack = new Stack<TreeNode>();
public BSTIterator(TreeNode root) {
pushAll(root);
}
/** @return whether we have a next smallest number */
public boolean hasNext() {
return !stack.isEmpty();
}
/** @return the next smallest number */
/* The average time complexity
* of next() function is O(1) indeed in your case. As the next function can
* be called n times at most, and the number of right nodes in
* self.pushAll(tmpNode.right) function is maximal n in a tree which has n
* nodes, so the amortized time complexity is O(1).
*/
public int next() {
TreeNode tmpNode = stack.pop();
pushAll(tmpNode.right);
return tmpNode.val;
}
private void pushAll(TreeNode node) {
while (node != null) {
stack.push(node);
node = node.left;
}
}
}