Update Binary Search Tree Iterator.py

Author: sugia

Binary Search Tree Iterator.py

@@ -3,10 +3,7 @@
 
 Calling next() will return the next smallest number in the BST.
 
-Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.
-
-Credits:
-Special thanks to @ts for adding this problem and creating all test cases.
+Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree. 
 '''
 
 # Definition for a  binary tree node
@@ -16,35 +13,40 @@
 #         self.left = None
 #         self.right = None
 
+# Definition for a  binary tree node
+# class TreeNode(object):
+#     def __init__(self, x):
+#         self.val = x
+#         self.left = None
+#         self.right = None
+
 class BSTIterator(object):
     def __init__(self, root):
         """
         :type root: TreeNode
         """
         self.stack = []
-        self.cur = root
-        
+        self.push(root)
 
+    def push(self, node):
+        while node:
+            self.stack.append(node)
+            node = node.left
+            
     def hasNext(self):
         """
         :rtype: bool
         """
-        return self.stack or self.cur
+        return len(self.stack) > 0
 
     def next(self):
         """
         :rtype: int
         """
-        while self.cur:
-            self.stack.append(self.cur)
-            self.cur = self.cur.left
-            
-        self.cur = self.stack.pop()
-        node = self.cur
-        self.cur = self.cur.right
-        
-        return node.val
-
+        top = self.stack.pop()
+        self.push(top.right)
+        return top.val
+    
 # Your BSTIterator will be called like this:
 # i, v = BSTIterator(root), []
 # while i.hasNext(): v.append(i.next())