173. Binary Search Tree Iterator

MediumStackTreeDesignBinary Search TreeBinary TreeIterator
Leetcode Link

Problem Description

The problem requires the creation of a class named BSTIterator that simulates an in-order iterator for a binary search tree (BST). An iterator, in this context, allows us to traverse the BST in a specific order - in-order traversal means that we visit the nodes in the increasing order of their values.

Here are the primary functionalities that need to be implemented:

  • BSTIterator(TreeNode root): The constructor of the class will take the root of the binary search tree and initiate the traversal. Essentially the constructor's job is to prepare the state of the iterator to begin the in-order traversal.

  • boolean hasNext(): This method should inform the user if there are more nodes to visit in the in-order traversal. It returns true if additional nodes are available; otherwise, it returns false.

  • int next(): This method is used to advance the iterator to the next number in the in-order traversal and return the value of the current number pointed to by the iterator.

By starting with a virtual pointer pointing to a non-existent smallest number, the iterator's first next() call will return the smallest number in the BST.

The problem guarantees that next() will only be called when it is valid to do so, meaning there will always be a next number to return.

Intuition

The natural way to perform an in-order traversal of a BST involves a recursive approach, visiting the left subtree, the root node, and then the right subtree. However, implementing an iterator with a purely recursive approach is not that straightforward. The BSTIterator should be able to pause the traversal between nodes and resume it when next() is called again.

To simulate this iterative process, a stack can be used. The stack can store nodes not yet visited, emulating the call stack that would be built up by a recursive in-order traversal. Upon initialization (__init__), the leftmost path of the BST is added to the stack, which composes all nodes that need to be visited before visiting the root (in order to respect the in-order traversal).

The next() method pops the next node from the stack (the current smallest element), processes it, and pushes all the left children of the right child, if any, onto the stack. This step ensures that after returning the current smallest number, the iterator is now pointing at the next smallest number.

The hasNext() method is straightforward, as all it needs to do is check if there are any nodes left in our stack. If the stack is not empty, the answer is true, otherwise, it is false.

By using a stack to store the next nodes to visit and by intelligently adding nodes to the stack while processing the next() method, we make sure to have the next value ready to be returned by next() while still using only O(h) memory, where h is the height of the tree. This is because at any given time, the stack only holds the nodes from the root to the leaf of the current path in the BST that we are visiting.

Learn more about Stack, Tree, Binary Search Tree and Binary Tree patterns.

Solution Approach

The implementation of the BSTIterator includes the use of a stack data structure. A stack follows the Last In, First Out (LIFO) principle, which is useful in simulating the in-order traversal iteratively.

In the constructor __init__, we initialize an empty stack. Then, starting from the given root node, we push onto the stack all the left children of the current node. We do so iteratively until we reach the leftmost node (which is the smallest element in the BST).

self.[stack](/problems/stack_intro) = []
while root:
    self.stack.append(root)
    root = root.left

This initial step ensures that calling next() for the first time will return the smallest element as per in-order traversal.

The next() method is responsible for popping the stack, which will give us the next element in the in-order traversal. Once we pop the stack, we must ensure that the next next() call will return the correct subsequent element. For this, we check if the popped node has a right child. If it does, we follow its leftmost path and push every node onto the stack, effectively doing the same operation we initially did for the root node.

cur = self.[stack](/problems/stack_intro).pop()
node = cur.right
while node:
    self.stack.append(node)
    node = node.left
return cur.val

The hasNext() method simply checks if there are more nodes to be visited in the stack.

return len(self.[stack](/problems/stack_intro)) > 0

It gives us a boolean indicating whether or not we still have a successor to visit.

Overall, the algorithm executes an in-order traversal on-demand, yielding the nodes one at a time while ensuring that the space complexity only depends on the height of the tree. This on-demand characteristic is achieved by combining the iterative traversal pattern with a stack that acts as a controlled version of the system call stack that would be used in a recursive approach.

Thus, with each call to next(), the iterator moves to the next node in the in-order sequence, while hasNext() provides a quick check to see if the traversal has been completed.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's consider a BST composed of the following nodes: 5, 3, 6, 2, 4, wherein 5 is the root. Its structure looks like this:

    5
   / \
  3   6
 / \
2   4

We would like to traverse this tree in an in-order manner using the BSTIterator class, which means we'll visit the nodes in the order 2, 3, 4, 5, 6.

First, we'll create an instance of BSTIterator with the root node (value 5):

iterator = BSTIterator(root)

During initialization within the constructor, the leftmost path of the BST is traversed, and all the nodes along this path are added onto the stack. After initialization, the stack would contain the nodes with the values [5, 3, 2], with 2 at the top since we pushed the leftmost nodes first (corresponding to the smallest elements).

Now, let's begin traversal with next() and hasNext() operations.

  1. Call next() for the first time. This should return 2, which is the smallest element. The stack now contains [5, 3].

    print(iterator.next())  # prints 2

  2. Call hasNext(), which will return True because the stack is not empty.

    print(iterator.hasNext())  # prints True

  3. Call next() again. It should return the next smallest element, which is 3. In the next() function, we also check if the node has a right child. Since 3 has a right child (4), we add it to the stack. The stack is now [5, 4].

    print(iterator.next())  # prints 3

  4. Call next() again. It will return 4 and since 4 has no right child, the stack remains [5].

    print(iterator.next())  # prints 4

  5. If we call hasNext() now, it should still return True since we have nodes yet to traverse.

    print(iterator.hasNext())  # prints True

  6. Call next() again and it returns 5. Since 5 has a right child which is 6, we push 6 onto the stack. The stack now contains only [6].

    print(iterator.next())  # prints 5

  7. Lastly, we call next() for the final node which is 6, and the stack becomes empty. There are no nodes left to traverse.

    print(iterator.next())  # prints 6

  8. One final call to hasNext() would now return False, signaling that the traversal is complete.

    print(iterator.hasNext())  # prints False

Throughout this process, the BSTIterator provided us with the next smallest number on each call to next() without requiring us to perform the in-order traversal in one go. The usage of a stack allowed us to track and control the traversal process.

Solution Implementation

1class TreeNode:
2    """Definition for a binary tree node."""  
3    def __init__(self, val=0, left=None, right=None):
4        """Initializes the tree node with a value, and optional left and right children."""
5        self.val = val
6        self.left = left
7        self.right = right
8
9class BSTIterator:
10    """Iterator over the nodes of a binary search tree (BST) that returns values in the ascending order."""
11  
12    def __init__(self, root: TreeNode):
13        """Initializes the iterator using the root of a BST."""
14        self.stack = []
15      
16        # Initialize the stack with all the nodes from root to the leftmost leaf.
17        while root:
18            self.stack.append(root)
19            root = root.left
20
21    def next(self) -> int:
22        """Returns the next smallest number in the BST."""
23        # The top element of the stack is the next smallest element.
24        current_node = self.stack.pop()
25      
26        # If there is a right subtree, push all the nodes from the right child 
27        # to the leftmost leaf onto the stack.
28        node = current_node.right
29        while node:
30            self.stack.append(node)
31            node = node.left
32          
33        # Return the value of the next smallest node.
34        return current_node.val
35
36    def hasNext(self) -> bool:
37        """Returns whether we have a next smallest number."""
38        # If the stack is not empty, there are still nodes available.
39        return len(self.stack) > 0
40
41# How to use the BSTIterator class:
42# obj = BSTIterator(root)
43# next_val = obj.next()
44# has_next = obj.hasNext()
45
1import java.util.Deque;
2import java.util.LinkedList;
3
4/**
5 * Definition for a binary tree node.
6 */
7class TreeNode {
8    int val;
9    TreeNode left;
10    TreeNode right;
11    TreeNode() {}
12    TreeNode(int val) { this.val = val; }
13    TreeNode(int val, TreeNode left, TreeNode right) {
14        this.val = val;
15        this.left = left;
16        this.right = right;
17    }
18}
19
20/**
21 * This class implements an iterator over a binary search tree (BST) using the "controlled" in-order traversal technique.
22 */
23class BSTIterator {
24    // Stack to store the path from root to the next smallest element.
25    private Deque<TreeNode> stack = new LinkedList<>();
26
27    /**
28     * Constructor that initializes the stack with the leftmost path of the BST (the smallest elements).
29     *
30     * @param rootNode the root node of the BST
31     */
32    public BSTIterator(TreeNode rootNode) {
33        // Start from the root and push all the leftmost nodes to the stack.
34        while (rootNode != null) {
35            stack.offerLast(rootNode);
36            rootNode = rootNode.left;
37        }
38    }
39
40    /**
41     * @return the next smallest number
42     */
43    public int next() {
44        // Pop the topmost node, which is the next smallest element in the BST.
45        TreeNode currentNode = stack.pollLast();
46      
47        // If the popped node has a right child, push all the leftmost nodes starting from the right child to the stack.
48        TreeNode node = currentNode.right;
49        while (node != null) {
50            stack.offerLast(node);
51            node = node.left;
52        }
53      
54        // Return the value of the next smallest node.
55        return currentNode.val;
56    }
57
58    /**
59     * @return whether we have a next smallest number
60     */
61    public boolean hasNext() {
62        // If the stack is not empty, there are still nodes to be visited.
63        return !stack.isEmpty();
64    }
65}
66
67/**
68 * Example of how the BSTIterator would be used:
69 * BSTIterator iterator = new BSTIterator(root);
70 * int val1 = iterator.next();    // Get the next smallest number
71 * boolean hasNxt = iterator.hasNext();    // Check if there is a next smallest number
72 */
73
1#include <stack>
2
3// Definition for a binary tree node.
4struct TreeNode {
5    int val;
6    TreeNode *left;
7    TreeNode *right;
8    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
9    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
10};
11
12class BSTIterator {
13private:
14    std::stack<TreeNode*> nodeStack; // Use std namespace and CamelCase for variables.
15
16public:
17    // Constructor initializes the iterator using the root of a binary search tree.
18    BSTIterator(TreeNode* root) {
19        // Go all the way to the left-most node, simulating an in-order traversal.
20        while (root != nullptr) {
21            nodeStack.push(root);
22            root = root->left;
23        }
24    }
25
26    // Returns the next smallest number in the BST.
27    int next() {
28        // The top element of the stack is the next smallest element.
29        TreeNode* currentNode = nodeStack.top();
30        nodeStack.pop();
31      
32        // If there is a right subtree, push all the left-most nodes of that subtree to the stack.
33        TreeNode* node = currentNode->right;
34        while (node != nullptr) {
35            nodeStack.push(node);
36            node = node->left;
37        }
38      
39        // Return the value of the next smallest element.
40        return currentNode->val;
41    }
42
43    // Returns whether we have a next smallest number.
44    bool hasNext() {
45        // If the stack is not empty, there is a next element.
46        return !nodeStack.empty();
47    }
48};
49
50/**
51 * The BSTIterator object will be instantiated and called as such:
52 * BSTIterator* obj = new BSTIterator(root);
53 * int param_1 = obj->next();
54 * bool param_2 = obj->hasNext();
55 */
56
1// Global stack to keep track of the nodes.
2let nodeStack: TreeNode[] = [];
3
4// The function that initializes the stack with the leftmost path of the tree.
5function initialize(root: TreeNode | null): void {
6    nodeStack = []; // Clearing any existing values in the global stack.
7    pushLeftBranch(root);
8}
9
10// The helper function to push the left branch of a node onto the stack.
11function pushLeftBranch(node: TreeNode | null): void {
12    while (node !== null) {
13        nodeStack.push(node);
14        node = node.left;
15    }
16}
17
18// The function that moves the iterator to the next element and returns its value.
19// This corresponds to the "next" method of the BSTIterator class.
20function getNext(): number {
21    if (nodeStack.length === 0) {
22        throw new Error("No next element.");
23    }
24
25    const currentNode = nodeStack.pop(); // Get the next node from the top of the stack.
26
27    pushLeftBranch(currentNode.right); // If there is a right child, push its leftmost path onto the stack.
28
29    return currentNode.val; // Return the current node's value.
30}
31
32// The function that checks if the iterator can move forward.
33// This corresponds to the "hasNext" method of the BSTIterator class.
34function hasNext(): boolean {
35    return nodeStack.length > 0; // If the stack is not empty, there is a next element.
36}
37
38// Example usage:
39// initialize(root); // Replace 'root' with the root of your binary search tree.
40// const value = getNext(); // Gets the next smallest number in the BST.
41// const canMoveNext = hasNext(); // Checks if there is a next number in the BST.
42

Time and Space Complexity

The provided code implements a BSTIterator for a binary search tree (BST). The __init__, next, and hasNext functions form the main interface of an iterator that provides a way to access BST elements in ascending order without having to store all of them at once.

  • __init__(self, root: TreeNode):

    • Time Complexity: The constructor has a time complexity of O(h), where h is the height of the tree. This is because the constructor iteratively goes down to the leftmost node, effectively traversing one side of the tree.
    • Space Complexity: The constructor has a space complexity of O(h) as well, due to the same iteration to the leftmost node. In the worst case, all nodes from the root node to the leftmost node are stored in the stack.

  • next(self) -> int:

    • Time Complexity: The average time complexity of next is O(1), since it amortizes over the number of calls. Each node is pushed to and popped from the stack exactly once. In the worst case of a single next call, the complexity can be O(h) when we traverse to the leftmost node of the right subtree.
    • Space Complexity: The space complexity does not change, which is still O(h). Any changes are temporary and only within the scope of the next call.

  • hasNext(self) -> bool:

    • Time Complexity: The time complexity is O(1) for checking if the stack is empty.
    • Space Complexity: There is no additional space used besides the existing stack, so the space complexity is O(1) for this operation.

In summary, while next operation averages O(1) time over the course of the entire traversal of the tree, certain calls can be O(h). The worst-case space complexity remains O(h) due to the stack size, which correlates to the height of the tree.

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

Which of the following is a good use case for backtracking?

Recommended Readings

Stack Intro

Stack Intro Imagine you have a pile of books on your desk If you want to add a new book you place it on top If you want to read a book you take it from the top And if you simply want to see which book is on the top you

Everything About Trees

Everything About Trees A tree is a type of graph data structure composed of nodes and edges Its main properties are It is acyclic doesn't contain any cycles There exists a path from the root to any node Has N 1 edges where N is the number of nodes in the tree and

Binary Search Tree Intro

Binary Search Tree Intro Definition A Binary Search Tree or BST is a rooted binary tree with the value of each internal node being greater than all the values in the respective node's left subtree and less than the ones in its right subtree An empty tree is a BST since it satisfies the

Want a Structured Path to Master System Design Too? Don’t Miss This!