For this problem, if don't consider the follow up limitation, the solution is very easy:
The time complexity is O(1), the space complexity is O(n), n is the total node number of the tree.
class BSTIterator { List<Integer> list = new ArrayList<>(); Iterator<Integer> it; public BSTIterator(TreeNode root) { inOrder(root); it = list.iterator(); } private void inOrder(TreeNode root){ if(root==null) return; inOrder(root.left); list.add(root.val); inOrder(root.right); } public int next() { return it.next(); } public boolean hasNext() { return it.hasNext(); } }
But the problem have a follow up:
Could you implement next()
and hasNext()
to run in average O(1)
time and use O(h)
memory, where h
is the height of the tree?
I can think of a solution with next()
and hasNext()
to run in average O(1)
time, and use O(1) memory, but this solution need to change the tree's structure. This solution conver the tree to a double direction linked list, inspired by: 426. Convert Binary Search Tree to Sorted Doubly Linked List
class BSTIterator { private TreeNode head, pre = new TreeNode(); public BSTIterator(TreeNode root) { head = root; if(head!=null){ while(head.left!=null) head = head.left; //1. find the smallest element inOrder(root); pre.right = head; head.left = pre; pre.right=head; } } private void inOrder(TreeNode root){ if(root == null) return; inOrder(root.left); pre.right = root; root.left = pre; pre = root; inOrder(root.right); } int count =0; public int next() { count++; pre=pre.right; return pre.val; } public boolean hasNext() { return count>0 && pre.right!=head||count==0; } }
The following is other's solution, the space complexity is O(h), but either next() or hasNext() cannot have O(1) time complexity:
class BSTIterator { private Stack<TreeNode> stk = new Stack<>(); public BSTIterator(TreeNode root) { pushLeft(root); } private pushLeft(TreeNode root) { while (root != null) { stk.push(root); root = root.left; } } public int next() { TreeNode temp = stk.pop(); pushLeft(temp.right) return temp.val; } public boolean hasNext() { return !stk.isEmpty(); } }