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.

Callingnext()will return the next smallest number in the BST.

Note:next()andhasNext()should run in average O(1) time and uses O(h) memory, wherehis the height of the tree.

Credits: Special thanks to @ts for adding this problem and creating all test cases.

Thoughts:

1. Using stack to record path.

2. When initialized, traverse to the leftmost

3. when retrieving, pop out from the stack and use its right child (if it does have) as a root node to put more nodes in front of the stack (In order traversal)

Code Python T:O(1) S:O(h)

# 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 = []
        cur = root
        while cur is not None:
            self.stack.append(cur)
            if cur.left:
                cur = cur.left
            else:
                break
        # print("len(self.stack) %s" %(len(self.stack)))
    def hasNext(self):
        """
        :rtype: bool
        """
        return self.stack

    def next(self):
        """
        :rtype: int
        """
        if not self.stack: return None

        ans = self.stack.pop() 
        cur = ans
        # check whether there is a right child, if there is, traverse towards the leftmost 
        # of the right child.
        if cur.right:
            cur = cur.right
            while cur:
                self.stack.append(cur)
                if cur.left:
                    cur = cur.left
                else: 
                    break

        return ans.val

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

from siyang2's post