Design an iterator over a binary search tree with the following properties:
- Elements are visited in ascending order (i.e. an inorder traversal)
- next() and hasNext() queries run in O(1) time in average.
Example
For the following binary search tree, inorder traversal by using iterator is [1, 6, 10, 11, 12]
10
/
1 11
6 12
Challenge
Extra memory usage O(h), h is the height of the tree.
Super Star: Extra memory usage O(1)
Solution: O(h) space.
1 /** 2 * Definition of TreeNode: 3 * public class TreeNode { 4 * public int val; 5 * public TreeNode left, right; 6 * public TreeNode(int val) { 7 * this.val = val; 8 * this.left = this.right = null; 9 * } 10 * } 11 * Example of iterate a tree: 12 * Solution iterator = new Solution(root); 13 * while (iterator.hasNext()) { 14 * TreeNode node = iterator.next(); 15 * do something for node 16 * } 17 */ 18 public class Solution { 19 private Stack<TreeNode> nodeStack; 20 21 //@param root: The root of binary tree. 22 public Solution(TreeNode root) { 23 nodeStack = new Stack<TreeNode>(); 24 //Initialize first, then determine null, otherwise, the hasNext() function will cause problem. 25 if (root==null) return; 26 nodeStack.push(root); 27 TreeNode cur = root.left; 28 while (cur!=null){ 29 nodeStack.push(cur); 30 cur = cur.left; 31 } 32 } 33 34 //@return: True if there has next node, or false 35 public boolean hasNext() { 36 if (nodeStack.isEmpty()) return false; 37 else return true; 38 } 39 40 //@return: return next node 41 public TreeNode next() { 42 if (nodeStack.isEmpty()) return null; 43 TreeNode next = nodeStack.pop(); 44 if (next.right==null) return next; 45 else { 46 nodeStack.push(next.right); 47 TreeNode cur = next.right.left; 48 while (cur!=null){ 49 nodeStack.push(cur); 50 cur = cur.left; 51 } 52 return next; 53 } 54 } 55 }