Implement Binary Search Tree(BST) Iterator

Binary Search Trees (BSTs) are data structures, in computer science. They are commonly used for searching, insertion and deletion operations. In situations it becomes necessary to traverse a BST in an order. For example an in order traversal visits nodes in ascending order based on their values. This article explores the creation of a BSTIterator class that facilitates the in order traversal of a search tree.

In this article, we will 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. The pointer should be initialized to a non-existent number smaller than any element in the BST.
• boolean hasNext(): Returns true if there exists a number in the traversal to the right of the pointer, otherwise returns false.
• int next(): Moves the pointer to the right, then returns the number at the pointer.

Notice that by initializing the pointer to a non-existent smallest number, the first call to the next() will return the smallest element in the BST. You may assume that the next() calls will always be valid. That is, there will be at least a next number in the in-order traversal when the next() is called.

Note: Time complexity for next, hasNext – O(1) and Space Complexity – O(N)

Approach and Implementation:

To achieve this functionality efficiently we can utilize a stack to keep track of the nodes along the path of the in-order traversal. Here is an explanation of how we can implement the BSTIterator class step by step;

1. Constructor (BSTIterator(TreeNode root)):

In the constructor we initialize the stack. Populate it with nodes from the leftmost path of the BST. This ensures that initially, the stack contains the elements of the tree.

Java

public BSTIterator(TreeNode root) {
    stack = new Stack<>();
    pushAllLeft(root);
}

The pushAllLeft function is a helper function that pushes all nodes from the node to its leftmost child onto the stack.

2. Boolean hasNext():

This method checks whether there are elements to traverse. It returns true if there are still elements remaining in the, in order traversal.

Java

public boolean hasNext() {
    return !stack.isEmpty();
}

3. Int next():

The next() function is responsible, for advancing the pointer to the following element during an in order traversal and returning that element. This is accomplished by removing the node from the stack and subsequently adding all the children of the right subtree of that node back, onto the stack.

Java

public int next() {
    TreeNode node = stack.pop();
    if (node.right != null) {
        pushAllLeft(node.right);
    }
    return node.val;
}

To guarantee that all left children of the subtree are included we also utilize the pushAllLeft function in this scenario.

Example:

Input: 7

/ \

3 15

/ \

9 20

Output: 3, 7, 9, 15, 20.

Below is the implementation of the BST Iterator:

C++

#include <iostream>
#include <stack>

using namespace std;

class TreeNode {
public:
    int val;
    TreeNode* left;
    TreeNode* right;

    TreeNode(int val) : val(val), left(nullptr), right(nullptr) {}
};

class BSTIterator {
private:
    stack<TreeNode*> st;

    void pushAllLeft(TreeNode* node) {
        while (node != nullptr) {
            st.push(node);
            node = node->left;
        }
    }

public:
    BSTIterator(TreeNode* root) {
        pushAllLeft(root);
    }

    bool hasNext() {
        return !st.empty();
    }

    int next() {
        TreeNode* node = st.top();
        st.pop();

        if (node->right != nullptr) {
            pushAllLeft(node->right);
        }

        return node->val;
    }
};

int main() {
    TreeNode* root = new TreeNode(7);
    root->left = new TreeNode(3);
    root->right = new TreeNode(15);
    root->right->left = new TreeNode(9);
    root->right->right = new TreeNode(20);

    BSTIterator iterator(root);

    while (iterator.hasNext()) {
        cout << iterator.next() << endl;
    }
    cout << endl;

    return 0;
}

Java

// Java code to implement
// Binary Search Tree iterator
import java.util.Stack;

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode(int val) { this.val = val; }
}

class BSTIterator {
    private Stack<TreeNode> stack;

    public BSTIterator(TreeNode root)
    {
        stack = new Stack<>();
        pushAllLeft(root);
    }

    public boolean hasNext() { return !stack.isEmpty(); }

    public int next()
    {
        TreeNode node = stack.pop();
        if (node.right != null) {
            pushAllLeft(node.right);
        }
        return node.val;
    }

    private void pushAllLeft(TreeNode node)
    {
        while (node != null) {
            stack.push(node);
            node = node.left;
        }
    }

    // Drivers code
    public static void main(String[] args)
    {
        TreeNode root = new TreeNode(7);
        root.left = new TreeNode(3);
        root.right = new TreeNode(15);
        root.right.left = new TreeNode(9);
        root.right.right = new TreeNode(20);

        BSTIterator iterator = new BSTIterator(root);

        while (iterator.hasNext()) {
            System.out.println(iterator.next());
        }
    }
}

Python3

# Python implementation
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

class BSTIterator:
    def __init__(self, root):
        self.st = []
        self.pushAllLeft(root)

    def pushAllLeft(self, node):
        while node:
            self.st.append(node)
            node = node.left

    def hasNext(self):
        return len(self.st) > 0

    def next(self):
        node = self.st.pop()

        if node.right:
            self.pushAllLeft(node.right)

        return node.val

root = TreeNode(7)
root.left = TreeNode(3)
root.right = TreeNode(15)
root.right.left = TreeNode(9)
root.right.right = TreeNode(20)

iterator = BSTIterator(root)

while iterator.hasNext():
    print(iterator.next())

# This code is contributed by Tapesh(tapeshdua420)

C#

// C# Implementation
using System;
using System.Collections.Generic;

public class TreeNode
{
    public int val;
    public TreeNode left;
    public TreeNode right;

    public TreeNode(int val)
    {
        this.val = val;
    }
}

public class BSTIterator
{
    private Stack<TreeNode> stack;

    public BSTIterator(TreeNode root)
    {
        stack = new Stack<TreeNode>();
        PushAllLeft(root);
    }

    public bool HasNext()
    {
        return stack.Count > 0;
    }

    public int Next()
    {
        TreeNode node = stack.Pop();
        if (node.right != null)
        {
            PushAllLeft(node.right);
        }
        return node.val;
    }

    private void PushAllLeft(TreeNode node)
    {
        while (node != null)
        {
            stack.Push(node);
            node = node.left;
        }
    }

    // Driver code
    public static void Main(string[] args)
    {
        TreeNode root = new TreeNode(7);
        root.left = new TreeNode(3);
        root.right = new TreeNode(15);
        root.right.left = new TreeNode(9);
        root.right.right = new TreeNode(20);

        BSTIterator iterator = new BSTIterator(root);

        while (iterator.HasNext())
        {
            Console.WriteLine(iterator.Next());
        }
    }
}

// This code is contributed by Tapesh (tapeshdua420)

Javascript

// Define a TreeNode class to represent nodes in the binary search tree
class TreeNode {
    constructor(val) {
        this.val = val;
        this.left = null;
        this.right = null;
    }
}

// Define a BSTIterator class
class BSTIterator {
    constructor(root) {
        this.stack = []; // Use a stack to keep track of nodes
        this.pushAllLeft(root); // Initialize the stack with leftmost nodes
    }

    // Helper method to push all left child nodes onto the stack
    pushAllLeft(node) {
        while (node !== null) {
            this.stack.push(node);
            node = node.left;
        }
    }

    // Check if there are more elements to iterate
    hasNext() {
        return this.stack.length > 0;
    }

    // Get the next element in ascending order
    next() {
        const node = this.stack.pop(); // Pop the top element from the stack

        if (node.right !== null) {
            this.pushAllLeft(node.right); // Push all left child nodes of the right subtree
        }

        return node.val; // Return the value of the popped node
    }
}

// Create a binary search tree
const root = new TreeNode(7);
root.left = new TreeNode(3);
root.right = new TreeNode(15);
root.right.left = new TreeNode(9);
root.right.right = new TreeNode(20);

// Initialize the BSTIterator with the root node
const iterator = new BSTIterator(root);

// Print the elements in ascending order using the iterator
while (iterator.hasNext()) {
    console.log(iterator.next());
}

Output

3
7
9
15
20

Time Complexity: O(N)
Auxiliary Space: O(N)