Friedrich Ewald My Personal Website

Leetcode: Binary Search Tree Iterator

Given a binary tree, write an iterator that performs in-order iteration and returns next as well as hasNext. It is guaranteed that next is never called on an empty tree. The naive approach looks as follows. Note that the space complexity here is O(n). 

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class BSTIterator:
  def __init__(self, root: Optional[TreeNode]):
    self._list = []
    self._buildTree(root)

  def _buildTree(self, root: Optional[TreeNode]):
    if root is None:
      return
    self._buildTree(root.left)
    self._list.append(root)
    self._buildTree(root.right)

  def next(self) -> int:
    return self._list.pop(0).val

  def hasNext(self) -> bool:
    return len(self._list) > 0
To further improve this, we can store the current nodes in a stack so that we only need to pop from the stack with time complexity O(1). The approach here is that we only iterate over the tree until we find the first leaf node, which happens on average after the height h of the tree. When popping the first element, we look for the next leaf node which can be found by going to the right and subsequently to the left. Repeat this process until the list is empty. 
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class BSTIterator:
  def __init__(self, root: Optional[TreeNode]):
    self._list = self._buildTree(root)

  def _buildTree(self, root: Optional[TreeNode]) -> List[TreeNode]:
    if root is None:
      return []
      
    return self._buildTree(root.left) + [root]


  def next(self) -> int:
    node = self._list.pop(0)
    self._list = self._buildTree(node.right) + self._list
    return node.val

  def hasNext(self) -> bool:
    return len(self._list) > 0

About the author

is a Staff Software Engineer with a Master's degree in Computer Science. He started this website in late 2015, mostly as a digital business card. He is interested in Go, Python, Ruby, SQL- and NoSQL-databases, machine learning and AI and is experienced in building scalable, distributed systems and micro-services at multiple larger and smaller companies.

Related Posts