Binary Search Tree Iterator

Intuition

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.

Solution


# 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()