Problem
Design an iterator over a binary search tree with the following rules:
Elements are visited in ascending order (i.e. an in-order traversal)
next() and hasNext() queries run in O(1) time in average.
Example
For the following binary search tree, in-order traversal by using iterator is [1, 6, 10, 11, 12]
10
/ \
1 11
\ \
6 12
思路
- 因为是中序遍历, 所以每次加一个节点, 需要把所有的左子节点push到stack里面
- hasNext()
- next()
- 取出stack顶元素, 并将该元素的右子树及右子树的所有左子节点加入到stack中
- 返回顶元素
- addToStack(curt)
- 将curt以及它的所有左子孙节点都加到stack中
public class BSTIterator {
Stack<TreeNode> stack;
private void addToStack(TreeNode root) {
while (root != null) {
stack.push(root);
root = root.left;
}
}
public BSTIterator(TreeNode root) {
stack = new Stack<>();
addToStack(root);
}
public boolean hasNext() {
if (!stack.isEmpty()) {
return true;
}
return false;
}
public TreeNode next() {
if (!hasNext()) {
return null;
}
TreeNode cur = stack.pop();
addToStack(cur.right);
return cur;
}
}