Binary Search Tree Iterator 86

Question

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)

Solution

利用一个stack实现。

next(): 因为时in-order，所以要讲当前root所有左孩子入栈，直到最左边为止。将最左边元素出栈，因为左边必然已经没有节点，所以将当前节点设为出栈元素的右孩子，再次重复前一步操作。

hasNext(): 当前root不为null（证明还有左边节点未入栈）或者栈不为空（还有元素未出栈）时为true。

代码如下：

public class BSTIterator {
    private Stack<TreeNode> stack = new Stack<>();
    private TreeNode curt;

    // @param root: The root of binary tree.
    public BSTIterator(TreeNode root) {
        curt = root;
    }

    //@return: True if there has next node, or false
    public boolean hasNext() {
        return (curt != null || !stack.isEmpty());
    }

    //@return: return next node
    public TreeNode next() {
        while (curt != null) {
            stack.push(curt);
            curt = curt.left;
        }

        curt = stack.pop();
        TreeNode node = curt;
        curt = curt.right;

        return node;
    }
}