LintCode-Implement Iterator of Binary Search Tree

 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 }