LeetCode/java/binary-search-tree-iterator.java

import java.util.ArrayDeque;

/* *
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */

class TreeNode {
  int val;
  TreeNode left;
  TreeNode right;

  TreeNode() {}

  TreeNode(int val) {
    this.val = val;
  }

  TreeNode(int val, TreeNode left, TreeNode right) {
    this.val = val;
    this.left = left;
    this.right = right;
  }
}

class BSTIterator {
  private enum Position {
    Todo,
    Done
  }

  private class Entry {
    public TreeNode node;
    public Position position;

    public Entry(TreeNode node, Position position) {
      this.node = node;
      this.position = position;
    }

    public void expand(ArrayDeque<Entry> stack) {
      if (node.right != null) {
        stack.addLast(new Entry(node.right, Position.Todo));
      }

      this.position = Position.Done;
      stack.addLast(this);

      if (node.left != null) {
        stack.addLast(new Entry(node.left, Position.Todo));
      }
    }
  }

  private ArrayDeque<Entry> stack = new ArrayDeque<>();

  public BSTIterator(TreeNode root) {
    stack.addLast(new Entry(root, Position.Todo));
  }

  public int next() {
    for (Entry last = stack.getLast();
        last != null && last.position == Position.Todo;
        last = stack.getLast()) {
      stack.removeLast().expand(stack);
    }

    return stack.removeLast().node.val;
  }

  public boolean hasNext() {
    Entry last = stack.peekLast();
    return last != null && (last.position != Position.Todo || last.node != null);
  }
}

/**
 * Your BSTIterator object will be instantiated and called as such: BSTIterator obj = new
 * BSTIterator(root); int param_1 = obj.next(); boolean param_2 = obj.hasNext();
 */