deepSWE.io

Binary Search Tree Iterator

TreesStack

Problem Statement

Implement the `BSTIterator` class that represents an iterator over the **in-order traversal** of a binary search tree (BST).

`BSTIterator(TreeNode root)` Initializes an object of the `BSTIterator` class. The `root` of the BST is given as part of the constructor.
`boolean hasNext()` Returns `true` if there is a next element in the in-order traversal, and `false` otherwise.
`int next()` Moves the iterator to the next element in the in-order traversal and returns the element's value.

Example

Operations: ["BSTIterator", "next", "next", "hasNext", "next", "hasNext", "next", "hasNext", "next", "hasNext"]

Output: [null, 3, 7, true, 9, true, 15, true, 20, false]

Algorithm Explanation

To implement the BST iterator, we need to perform an in-order traversal. A key insight is to use a stack to store the nodes we need to visit. This allows us to pause and resume the traversal.

Algorithm Steps

Constructor: When the iterator is initialized, we don't populate the stack with all nodes at once. Instead, we'll add nodes as we need them. We start by pushing all the left children of the root onto the stack.
`next()`: When `next()` is called, we pop a node from the stack. This is the smallest node we haven't visited yet. Before returning its value, we check if it has a right child. If it does, we push that right child and all of its left descendants onto the stack.
`hasNext()`: The `hasNext()` method is simple: if the stack is not empty, it means there are still nodes to visit, so we return `true`. Otherwise, we return `false`.

Constructor

Initializing BSTIterator.

Result: [ ]

Stack

BSTIterator Solution


class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.stack = []
        self._push_left(root)

    def next(self) -> int:
        # The node at the top of the stack is the next smallest element
        node = self.stack.pop()
        # If the node has a right child, push all its left descendants
        self._push_left(node.right)
        return node.val

    def hasNext(self) -> bool:
        # If the stack is not empty, there are more elements
        return len(self.stack) > 0

    def _push_left(self, node: Optional[TreeNode]):
        # Push all left children of a node onto the stack
        while node:
            self.stack.append(node)
            node = node.left