diff --git a/problems/binary-search-tree-iterator.java b/problems/binary-search-tree-iterator.java new file mode 100644 index 0000000..9b945f5 --- /dev/null +++ b/problems/binary-search-tree-iterator.java @@ -0,0 +1,88 @@ +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 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 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(); + */