Binary Search Tree Iterator 解答

Question

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.

Solution

When a problem relates to BST and sequence, we should think about in-order traversal of BST.

We find that the original codes of Binary Tree Inorder Traversal can be modified to fit following codes.

while (it.hasNext()) {
    System.out.println(it.next());        
}

Note that before print, we need to push nodes to stack first.

Inorder traversal time complexity is O(n), so for each next() step, average time is O(1).

And the stack costs O(h) because we either goes down or pop a node and then goes down.

/**
 * Definition for binary tree
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */

public class BSTIterator {
    private Stack<TreeNode> stack;
    private TreeNode current;

    public BSTIterator(TreeNode root) {
        stack = new Stack<TreeNode>();
        current = root;
        while (current != null) {
            stack.push(current);
            current = current.left;
        }
        
    }

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

    /** @return the next smallest number */
    public int next() {
        while (current != null) {
            stack.push(current);
            current = current.left;
        }
        TreeNode tmp = stack.pop();
        int result = tmp.val;
        current = tmp.right;
        return result;
    }
}

/**
 * Your BSTIterator will be called like this:
 * BSTIterator i = new BSTIterator(root);
 * while (i.hasNext()) v[f()] = i.next();
 */