173. Binary Search Tree Iterator

Leetcode Design Stack Tree

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.

Calling next() will return the next smallest number in the BST.

Note: next() and hasNext() should run in average $O(1)$ time and uses $O(h)$ memory, where $h$ is the height of the tree.

分析

这道题目看到的时候就可以想到用中序遍历，把中序遍历的结果放在List中，然后调用List.iterator()方法。但是这道题目的要求更高，它的空间复杂度要求是$O(h)$，其中$h$是树的高度。这样一来这种方法就不行了。仔细观察二叉树的中序遍历：(代码来自LeeCode 94. Binary Tree Inorder Traversal)

public List<Integer> inorderTraversal(TreeNode root) {
    List<Integer> res = new ArrayList<Integer>();
    Stack<TreeNode> stack = new Stack<TreeNode>();
    TreeNode cur = root;
    while (cur != null || !stack.isEmpty()) { 
        while (cur != null) { 
            // Travel to each node's left child,
            // till reach the left leaf
            stack.push(cur);
            cur = cur.left;             
        }        
        cur = stack.pop(); // Backtrack to higher level node A
        res.add(cur.val);  // Add the node to the result list
        cur = cur.right;   // Switch to A'right branch
    }
    return res;
}  

恰好每一个while循环有一个节点加入List，那么一个猜想就是我们可不可以把每个while循环拆分开，相当于一个每个while循环就是一个next()方法调用。那么怎么判断hasNext()呢，一定和while循环的条件差不多。

    private  Stack<TreeNode> stack;
    private TreeNode cur;

    public BSTIterator(TreeNode root) {
        stack = new Stack<TreeNode>();
        cur = root;
    }

    /** @return whether we have a next smallest number */
    public boolean hasNext() {
        return cur != null || stack.size() != 0;
    }

    /** @return the next smallest number */
    public int next() {
        while (cur != null) { 
            // Travel to each node's left child,
            // till reach the left leaf
            stack.push(cur);
            cur = cur.left;             
        }        
        cur = stack.pop(); // Backtrack to higher level node A
        int res = cur.val;  // Add the node to the result 
        cur = cur.right;   // Switch to A'right branch
        return res;
    }
}

代码和中序遍历基本一致，只改变了while循环为next()。