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
Challenge
Extra memory usage O(h), h is the height of the tree.
Super Star: Extra memory usage O(1)
Note
建立一个堆栈,先将最左边的结点从大到小压入栈,这样的话,为了实现迭代
即返回下一个node的next()
函数就要考虑右边的结点。比如example中的BST,stack第一次pop出来的结点是1
,然后就应该把它的右子树结点6
压入stack;又如pop出了root结点10
以后,就要把11
压入堆栈,这样,在pop出11
之后,再将12
压入堆栈。如此,实现next()函数。
Solution
public class BSTIterator {
//@param root: The root of binary tree.
Stack<TreeNode> stack;
public BSTIterator(TreeNode root) {
stack = new Stack<TreeNode>();
while (root != null) {
stack.push(root);
root = root.left;
}
}
//@return: True if there has next node, or false
public boolean hasNext() {
return !stack.isEmpty();
}
//@return: return next node
public TreeNode next() {
TreeNode cur = stack.pop();
TreeNode temp = cur;
if (cur.right != null) {
cur = cur.right;
while (cur != null) {
stack.push(cur);
cur = cur.left;
}
}
return temp;
}
}