Binary Search Tree iterator II

XXX. BST Iterator with Backward and Forward Navigation

Stack

In-Order Traversal

Problem Statement:

Design an iterator over a binary search tree (BST) that supports forward and backward navigation through the tree. The iterator must return values in ascending order.

Algorithm:

1. Use a stack to simulate in-order traversal of the BST. This stack will help with efficiently finding the next smallest node in the tree.
2. Maintain a list to store all visited nodes for backward navigation. Use a pointer to keep track of the current position in this list.
3. For the next() operation:
   • Check if there are unvisited nodes left in the stack.
   • If the pointer is already at the end of the visited nodes, compute the next node by processing the current node's right subtree.
4. For the prev() operation:
   • Check if the pointer can move backward in the visited list.

Complexity:

Time Complexity:

• next(): Average O(1), amortized over all nodes in the tree.
• prev(): O(1), since it simply accesses the previous node in the list.

Space Complexity: O(h + n), where h is the height of the tree (for the stack), and n is the number of nodes (for the visited list).

Java Implementation:

class BSTIterator {

    Stack stack;
    List arr;
    int pointer;

    public BSTIterator(TreeNode root) {
        stack = new Stack<>();
        arr = new ArrayList<>();
        pointer = -1;

        // Initialize the stack with the leftmost path
        while (root != null) {
            stack.push(root);
            root = root.left;
        }
    }

    public boolean hasNext() {
        return pointer < arr.size() - 1 || !stack.isEmpty();
    }

    public int next() {
        pointer++;

        // If the next element is not precomputed
        if (pointer == arr.size()) {
            TreeNode node = stack.pop();
            TreeNode next = node.right;
            while (next != null) {
                stack.push(next);
                next = next.left;
            }

            arr.add(node.val); // Add the current node value to the list
        }

        return arr.get(pointer);
    }

    public boolean hasPrev() {
        return pointer > 0;
    }

    public int prev() {
        pointer--;
        return arr.get(pointer);
    }
}
                

Python Implementation:

class BSTIterator:

    def __init__(self, root):
        self.stack = []
        self.arr = []
        self.pointer = -1

        # Initialize the stack with the leftmost path
        while root:
            self.stack.append(root)
            root = root.left

    def hasNext(self):
        return self.pointer < len(self.arr) - 1 or self.stack

    def next(self):
        self.pointer += 1

        if self.pointer == len(self.arr):
            node = self.stack.pop()
            next_node = node.right
            while next_node:
                self.stack.append(next_node)
                next_node = next_node.left

            self.arr.append(node.val)  # Add the current node value to the list

        return self.arr[self.pointer]

    def hasPrev(self):
        return self.pointer > 0

    def prev(self):
        self.pointer -= 1
        return self.arr[self.pointer]
                

C++ Implementation:

#include 
#include 
using namespace std;

class BSTIterator {
    stack stack;
    vector arr;
    int pointer;

public:
    BSTIterator(TreeNode* root) : pointer(-1) {
        // Initialize the stack with the leftmost path
        while (root) {
            stack.push(root);
            root = root->left;
        }
    }

    bool hasNext() {
        return pointer < (int)arr.size() - 1 || !stack.empty();
    }

    int next() {
        pointer++;

        if (pointer == arr.size()) {
            TreeNode* node = stack.top();
            stack.pop();
            TreeNode* next = node->right;
            while (next) {
                stack.push(next);
                next = next->left;
            }

            arr.push_back(node->val); // Add the current node value to the list
        }

        return arr[pointer];
    }

    bool hasPrev() {
        return pointer > 0;
    }

    int prev() {
        pointer--;
        return arr[pointer];
    }
};