from array204 import *
class _BSTIterator :      
  def __init__( self, root ):
     # Figure out how many nodes in the tree
    size = self._count_nodes( root )

     # Creates the array and fills it with the keys.
    self._theKeys = Array( size )
    self._curItem = 0    # Keep track of the next location in the array.
    self._bstTraversal( root )     
    self._curItem = 0    # Reset the current item index.
    
  def __iter__( self ):
    return self

   # Returns the next key from the array of keys 
  def __next__( self ):
    if self._curItem < len( self._theKeys ) :
      key = self._theKeys[ self._curItem ]
      self._curItem += 1
      return key
    else :
      raise StopIteration
      
   # Performs an inorder traversal used to build the array of keys. 
  def _bstTraversal( self, subtree ):
    if subtree is not None :
      self._bstTraversal( subtree.left )
      self._theKeys[ self._curItem ] = subtree.key
      self._curItem += 1
      self._bstTraversal( subtree.right )       
    
    # Count the number of nodes in this BST
  def _count_nodes( self, root ):
    if root == None:
      return 0
    else:
      l_count = self._count_nodes( root.left )
      r_count = self._count_nodes( root.right )
      return 1 + l_count + r_count