Friday, January 2, 2015

Binary Inorder Traversal Iterator

Problem

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.


For a binary search tree, kept returning the next smallest element is inorder traversal. Wikipedia has some great picture to show how to traversal in "preorder", "inorder" and "postorder". If you are not familiar with this, just go to Wikipedia.

The LeetCode always encourage us use loop other than recursion to solve tree traversal, because recursion may cost more extra space time for program stack and easy to understand.

When I see the problem, my first idea is push all the elements in the tree to a stack. But this cost O(n) memory. So a idea result is pop one from stack and push one to stack. Most method didn't do this precisely, just an amortized O(1) time and O(h) space. If the tree is pretty unbalanced, it is as same as push all the elements in stack. Here is code: private Stack nodes = new Stack<>();

public BSTIterator(TreeNode root) {
    pushToStack(root);
}

/** @return whether we have a next smallest number */
public boolean hasNext() {
    return !nodes.isEmpty();
}

/** @return the next smallest number */
public int next() {
    TreeNode tmp = nodes.pop();
    pushToStack(tmp.right);
    return tmp.val;
}

private void pushToStack(TreeNode node) {
    while (node != null) {
        nodes.push(node);
        node = node.left;
    }
}

The idea is keep going left until no more elements in left, than to right element and keep going left again. Redo above until no elements in the stack. The problems need to return elements inorder is the same idea. Binary Tree Inorder Tracersal could be solve use the same code structure:

public List<Integer> inorderTraversal(TreeNode root) {
    Stack<TreeNode> nodes = new Stack<>();
    ArrayList<Integer> result = new ArrayList<>();
    if (root == null)
        return result;

    pushToStack(root, nodes);

    while (!nodes.isEmpty()) {
        TreeNode tmp = nodes.pop();
        pushToStack(tmp.right, nodes);
        result.add(tmp.val);
    }
    return result;
}

public void pushToStack(TreeNode node, Stack<TreeNode> nodes) {
    while (node != null) {
        nodes.push(node);
        node = node.left;
    }
}

Written with StackEdit.

No comments:

Post a Comment