BST iterator

Hero Image

DT

Dhaval Trivedi

Co-founder, Airtribe

Problem Overview

Title: Binary Search Tree Iterator

The Binary Search Tree (BST) Iterator problem provides a Binary Search Tree and asks you to implement an iterator over the tree. The iterator should be able to return the elements of the BST in ascending order. Essentially, you are required to implement two functions:

  • next(): Returns the next smallest number in the BST.
  • hasNext(): Returns whether the next smallest number is available or not.

Understanding the Problem

The main goal is to mimic an iterator which efficiently retrieves elements in in-order traversal (left-root-right) from a BST. The challenge is to implement this with controlled space complexity and fast access to the next element, ideally in average O(1) time for next() and O(h) space where h is the height of the tree.

Example:

Consider a BST:

    7
   / \
  3  15
    /  \
   9   20
  1. BSTIterator initialized with root node of the BST.
  2. next() returns 3, hasNext() returns true
  3. next() returns 7, hasNext() returns true
  4. next() returns 9, hasNext() returns true
  5. next() returns 15, hasNext() returns true
  6. next() returns 20, hasNext() returns false

Key DSA Concepts and Theory

Binary Search Tree (BST)

A Binary Search Tree is a binary tree in which each node has a comparable key (or value) and satisfies the condition that the key in each node is greater than the keys in its left subtree and less than the keys in its right subtree.

In-order Traversal

In-order traversal of a BST retrieves the nodes in ascending order. During this traversal, we visit the left subtree, then the current node, and finally the right subtree.

Stack Data Structure

A stack is an abstract data type that serves as a collection of elements, with two principal operations:

  • Push: Adding an element to the collection.
  • Pop: Removing the most recently added element that was not yet removed.

In this problem, a stack is used to simulate the recursive in-order traversal iteratively.

Solution Approach

The key to solving this problem is to utilize a stack to perform an in-order traversal of the BST iteratively. Here's a breakdown of the approach:

  1. Initialization

    • Initialize a stack to keep track of the nodes.
    • In the constructor, push all the leftmost nodes of the tree onto the stack.
  2. hasNext() Method

    • Simply checks if there are any elements left in the stack.
  3. next() Method

    • Pop the top element from the stack, which is the next smallest element.
    • If the popped node has a right child, push all the leftmost nodes of this right subtree to the stack.

C++ Solution

class BSTIterator {
public:
    stack<TreeNode*> nodeStack;
    
    BSTIterator(TreeNode* root) {
        pushAllLeft(root);
    }
    
    /** @return the next smallest number */
    int next() {
        TreeNode *node = nodeStack.top();
        nodeStack.pop();
        // If the node has a right child, process its leftmost nodes
        if (node->right) {
            pushAllLeft(node->right);
        }
        return node->val;
    }
    
    /** @return whether we have a next smallest number */
    bool hasNext() {
        return !nodeStack.empty();
    }
    
private:
    void pushAllLeft(TreeNode* node) {
        while (node != nullptr) {
            nodeStack.push(node);
            node = node->left;
        }
    }
};

Java Solution

class BSTIterator {
    private Stack<TreeNode> stack;

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

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

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

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

Time and Space Complexity Analysis

Time Complexity

  • next(): Average O(1) because each node is processed exactly once and each call either pops or pushes a node.
  • hasNext(): O(1).

Space Complexity

  • O(h) where h is the height of the tree. This space is used by the stack to store nodes.

Common Mistakes to Avoid

  • Not handling empty trees correctly.
  • Forgetting to push left nodes of the right child in the next() method.
  • Misunderstanding the structure of the iterative in-order traversal.

Similar Problems on LeetCode

  1. Flatten Binary Tree to Linked List - LeetCode 114
  2. Inorder Successor in BST - LeetCode 285
  3. Kth Smallest Element in a BST - LeetCode 230

Additional Resources and References

This comprehensive coverage should give you a clear understanding of the solution approach and help you implement it in your own way. Happy coding!