We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Binary Search Tree Iterator

Number: 173

Difficulty: Medium

Paid? No

Companies: Meta, Microsoft, Amazon, Uber, Adobe, Media.net, Bloomberg, Apple, Google, LinkedIn

Problem Description

Design an iterator for a binary search tree (BST) that returns nodes during an in-order traversal. Implement the BSTIterator class with the following functionalities:

  • BSTIterator(TreeNode root): Initializes the iterator. The pointer is set to a value smaller than the smallest element.
  • boolean hasNext(): Returns true if there are further nodes in the traversal.
  • int next(): Advances the iterator to the next node and returns its value.

Key Insights

  • Use an iterative in-order traversal strategy with a stack.
  • In-order traversal for BST yields sorted order.
  • Pre-load the stack with all left descendants from the root.
  • After returning a node, process its right subtree by pushing its left branch.
  • This method maintains an average time complexity of O(1) per operation and uses O(h) memory, where h is the tree height.

Space and Time Complexity

Time Complexity: O(1) on average per operation (amortized) Space Complexity: O(h), where h is the height of the tree

Solution

We use a stack to simulate an in-order traversal. The process involves initially pushing all left-child nodes from the root into the stack. For each call to next(), we pop the top element (the current smallest node) and then push all the left descendents of the node's right child. This ensures that the next smallest element is always at the top of the stack. The hasNext() function then simply checks if the stack is non-empty.

Code Solutions

# 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):
        # Initialize an empty stack to simulate in-order traversal
        self.stack = []
        # Push all left nodes starting from the root
        self._push_left_branch(root)
    
    def _push_left_branch(self, node):
        # Push node and all its left children to the stack
        while node:
            self.stack.append(node)
            node = node.left
    
    def next(self):
        # Pop the top node which represents the next smallest element
        node = self.stack.pop()
        # If the node has a right child, push its left branch onto the stack
        if node.right:
            self._push_left_branch(node.right)
        # Return the node's value
        return node.val
    
    def hasNext(self):
        # The iterator has a next element if the stack is not empty
        return len(self.stack) > 0
← Back to All Questions