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.
对的
代码:
private Stack<TreeNode> stack;
private TreeNode root;
public BSTIterator(TreeNode root) {
this.root = root;
stack = new Stack<>();
pushToStack(root);
}
private void pushToStack(TreeNode node){
if(node == null) return;
TreeNode p = node;
while(p != null){
stack.push(p);
p = p.left;
}
}
/** @return whether we have a next smallest number */
public boolean hasNext() {
if(stack.isEmpty()) return false;
return true;
}
/** @return the next smallest number */
public int next() {
if(stack.isEmpty()) return -1;
TreeNode tree = stack.pop();
if(tree.right != null){
pushToStack(tree.right);
}
return tree.val;
}