Question Link:

https://leetcode.com/problems/binary-search-tree-iterator/


Solution

This problem sounds complicated, but it is not. The solution is just divided the algorithm of [LeetCode 144] Binary Tree Preorder Traversal into multiple steps of hasNext() and next() functions.

We keep the stack which is used in the inorder traverse. We only pop up when the next() function is called.

Python implementation

The python code is as follows.

[-] Python code accepted by LeetCode OJ
# 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
        """
        # Create the stack until we found the first node
        self.stack = []
        while root:
            self.stack.append(root)
            root = root.left

    def hasNext(self):
        """
        :rtype: bool
        """
        return len(self.stack) > 0

    def next(self):
        """
        :rtype: int
        """
        # Pop the next node from the stack
        ret = self.stack.pop()

        # Continue the inorder traverse
        N = ret.right
        while N:
            self.stack.append(N)
            N = N.left

        # Return
        return ret.val

# Your BSTIterator will be called like this:
# i, v = BSTIterator(root), []
# while i.hasNext(): v.append(i.next())