Lintcode 86. Binary Search Tree Iterator (Hard) (Python)

Binary Search Tree Iterator

Description:

Design an iterator over a binary search tree with the following rules:

Elements are visited in ascending order (i.e. an in-order traversal)
next() and hasNext() queries run in O(1) time in average.
Example
For the following binary search tree, in-order traversal by using iterator is [1, 6, 10, 11, 12]

10
/
1 11
\
6 12
Challenge
Extra memory usage O(h), h is the height of the tree.

Super Star: Extra memory usage O(1)

Code:

"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None

Example of iterate a tree:
iterator = BSTIterator(root)
while iterator.hasNext():
    node = iterator.next()
    do something for node 
"""


class BSTIterator:
    """
    @param: root: The root of binary tree.
    """
    def __init__(self, root):
        # do intialization if necessary
        self.curr = root
        self.stack = []
    """
    @return: True if there has next node, or false
    """
    def hasNext(self, ):
        # write your code here
        while self.curr:
            self.stack.append(self.curr)
            self.curr = self.curr.left
        return self.stack
    """
    @return: return next node
    """
    def next(self, ):
        # write your code here
        if self.hasNext():
            node = self.stack.pop()
            if node.right:
                self.curr = node.right
            return node
        else:
            return None