[LintCode] Binary Search Tree Iterator

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;
    }
}