LintCode-Implement Iterator of Binary Search Tree

Design an iterator over a binary search tree with the following properties:

1. Elements are visited in ascending order (i.e. an inorder traversal)
2. next() and hasNext() queries run in O(1) time in average.

Example

For the following binary search tree, inorder 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: O(h) space.

/**
 * Definition of TreeNode:
 * public class TreeNode {
 *     public int val;
 *     public TreeNode left, right;
 *     public TreeNode(int val) {
 *         this.val = val;
 *         this.left = this.right = null;
 *     }
 * }
 * Example of iterate a tree:
 * Solution iterator = new Solution(root);
 * while (iterator.hasNext()) {
 *    TreeNode node = iterator.next();
 *    do something for node
 * }
 */
public class Solution {
    private Stack<TreeNode> nodeStack;

    //@param root: The root of binary tree.
    public Solution(TreeNode root) {
        nodeStack = new Stack<TreeNode>();
        //Initialize first, then determine null, otherwise, the hasNext() function will cause problem.
        if (root==null) return;
        nodeStack.push(root);
        TreeNode cur = root.left;
        while (cur!=null){
            nodeStack.push(cur);
            cur = cur.left;
        }
    }

    //@return: True if there has next node, or false
    public boolean hasNext() {
        if (nodeStack.isEmpty()) return false;
        else return true;
    }
    
    //@return: return next node
    public TreeNode next() {
        if (nodeStack.isEmpty()) return null;
        TreeNode next = nodeStack.pop();
        if (next.right==null) return next;
        else {
            nodeStack.push(next.right);
            TreeNode cur = next.right.left;
            while (cur!=null){
                nodeStack.push(cur);
                cur = cur.left;
            }
            return next;
        }
    }
}