We could use stack to save the state of the recursion and iterate through the nodes one by one.
There are 2 different state, the state 0 is the node just putted into the stack and none of its sub-nodes are visited. The state 1 is that the left of this node has been visited, now we can visit itself and then put its right to the queue.
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class BSTIterator:
def __init__(self, root: TreeNode):
self.stack = [[root, 0]]
self._next = None
self._get_next()
def _get_next(self):
while self.stack:
node, state = self.stack.pop()
if node is None:
continue
if state == 0:
self.stack.append([node, 1])
self.stack.append([node.left, 0])
else:
self._next = node.val
self.stack.append([node.right, 0])
return
self._next = None
def next(self) -> int:
"""
@return the next smallest number
"""
tmp = self._next
self._get_next()
return tmp
def hasNext(self) -> bool:
"""
@return whether we have a next smallest number
"""
return self._next is not None
# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()