[LintCode] Binary Search Tree Iterator

421 查看

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

建立一个堆栈,先将最左边的结点从大到小压入栈,这样的话,为了实现迭代返回下一个nodenext()函数就要考虑右边的结点。比如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;
    }
}