Wednesday, January 18, 2017

Binary Search Tree Iterator

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.
/**
 * Your BSTIterator will be called like this:
 * BSTIterator i = new BSTIterator(root);
 * while (i.hasNext()) v[f()] = i.next();
 */

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

/*
Every time, we call next(), we will traverse to find the left most TreeNode of the current node if existed. If the current node is null, then we pop out a node from the stack and return its value. The first time we call next(), it takes O(h) time, but the other TreeNode in the stack will take only O(1) time when call the next(). And there are h TreeNode in the stack. So the average time complexity will be O(1) and space complexity will be O(h).

hasNext() and next(): time:O(1), space:O(h)
*/
public class BSTIterator {
    public Stack<TreeNode> stack = new Stack<>();
    public TreeNode root;

    public BSTIterator(TreeNode root) {
        this.root = root;
    }

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

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

No comments:

Post a Comment