Binary Search Tree (BST) Iterator Editorial

Practice Problem Link: Binary Search Tree (BST) Iterator

Please make sure to try solving the problem yourself before looking at the editorial.

Problem Statement

Implement the BSTIterator class. The BSTIterator will be used to iterate over the inorder traversal of a BST.

• BSTIterator(root) initializes an object of the BSTIterator class.
• boolean hasNext() returns whether there is a next element in the traversal.
• Node next() moves the pointer to the next node in the traversal and returns the node.

Assume that next() will be called only if there is a next node.

The iterator would be called like this depending on your language:

BSTIterator iterator = new BSTIterator(root);
while iterator.hasNext() {
    print iterator.next()
}

Approach

The idea is to perform a DFS like traversal.

• Initially store all the left nodes of the given BST in an array so that the last node will be the smallest value in the BST.
• To check whether there is a next element or not in the traversal, simply check if the array is empty or not.
• To move the pointer to the next node, remove the last node from the array. Now add all the left nodes of the right subtree of the removed topmost node in the array so that the new last node in the array will point to the next node in the inorder traversal. Now return the removed node.

Analysis

• Time Complexity: O(n)
• Auxiliary Space Complexity: O(Height of BST)

Implementation

C++

/* This is the Node class definition

class Node {
public:
    Node* left;
    Node* right;
    int data;

    Node(int data) {
        this->left = NULL;
        this->right = NULL;
        this->data = data;
    }
};
*/

class BSTIterator {
public:
	vector<Node*> treeNodes;
	void insertLeftNodes(Node* root) {
		while (root != NULL) {
			treeNodes.push_back(root);
			root = root->left;
		}
	}
    BSTIterator(Node* root) {
		insertLeftNodes(root);
    }
    
    bool hasNext() {
		return !treeNodes.empty();
    }
    
    Node* next() {
		int last = treeNodes.size() - 1;
		Node* currNode = treeNodes[last];
		treeNodes.pop_back();
		insertLeftNodes(currNode->right);
		return currNode;
    }
};

Java

/* This is the Node class definition

class Node {
	public Node left;
	public Node right;
	public int data;

	public Node(int data) {
		this.data = data;
	}
}
*/
class BSTIterator {
	List<Node> treeNodes = new ArrayList<Node>();
	void insertLeftNodes(Node root) {
		while (root != null) {
			treeNodes.add(root);
			root = root.left;
		}
	}
    public BSTIterator(Node root) {
		insertLeftNodes(root);
    }
    
    public boolean hasNext() {
		return !treeNodes.isEmpty();
    }
    
    public Node next() {
		int last = treeNodes.size() - 1;
		Node currNode = treeNodes.get(last);
		treeNodes.remove(last);
		insertLeftNodes(currNode.right);
		return currNode;
    }
}