173. Binary Search Tree Iterator

Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.

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.

Stack (Using stack to find BST next smallest is similar to in-order traverse the BST. ):

# 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 = []
        while root:
            self.stack.append(root)
            root = root.left
        
    def push2LeftMost(self, root):
        while(root):
            self.stack.append(root)
            root = root.left
            

    def hasNext(self):
        """
        :rtype: bool
        """
        if self.stack:
            return True
        else:
            return False
        

    def next(self):
        """
        :rtype: int
        """
        if self.hasNext():
            rst = self.stack.pop()
            # 这里就跟 inorder traverse 一样样的,非常关键;虽然 rst 的结果已经得到了,
            # 但是还是需要把 inorder traverse 的剩余情况 push 进 stack 内。即,一旦从 stack
            # 内 pop 出一个 node,它的非 None 右子节点以及非 None 右子节点的所有左支都要 push 进
            # stack。
            child = rst.right  
            while(child):
                self.stack.append(child)
                child = child.left
            return rst.val
        else:
            return None

Last updated