[LeetCode] 173. Binary Search Tree Iterator

Posted on | In

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.

Example:

173. Binary Search Tree Iterator

1
2
3
4
5
6
7
8
9
10
BSTIterator iterator = new BSTIterator(root);
iterator.next();    // return 3
iterator.next();    // return 7
iterator.hasNext(); // return true
iterator.next();    // return 9
iterator.hasNext(); // return true
iterator.next();    // return 15
iterator.hasNext(); // return true
iterator.next();    // return 20
iterator.hasNext(); // return false

Note:

  • next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.
  • You may assume that next() call will always be valid, that is, there will be at least a next smallest number in the BST when next() is called.

Explanation

  1. In the constructor, we can use in-order traversal iterative way to store the left child elements into a stack.

  2. In the hasNext function, we can check if stack is empty. If not empty, then return true.

  3. In the next function, we can return the pop node element. Also, we check if that pop node has right child, if it has right child, then we think of that right child is the root, and we push that “root” and all its left nodes to the stack like we did in the step 1.

Solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class BSTIterator {
    Stack<TreeNode> st;

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

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

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