LeetCode 1586 - Binary Search Tree Iterator II

LeetCode Problem 1586

Difficulty: 🟡 Medium
Topics: Stack, Tree, Design, Binary Search Tree, Binary Tree, Iterator

Solution

Problem Understanding

The problem requires implementing a BSTIterator class that allows forward and backward traversal over the in-order sequence of a Binary Search Tree (BST). In-order traversal of a BST visits nodes in ascending order. The iterator must support four operations: next(), prev(), hasNext(), and hasPrev(). Initially, the iterator is positioned just before the first element in the traversal, so the first call to next() returns the smallest element. The input is a BST represented by its root node, and the output is the sequence of elements accessed through the iterator methods.

The constraints indicate that the BST can contain up to 10^5 nodes and that there can be up to 10^5 calls to the iterator methods. The values of the nodes range from 0 to 10^6. This implies that we need a solution that is efficient both in time and space. Edge cases include a BST with a single node, calls to prev() immediately after initialization, and skewed trees (completely unbalanced).

Approaches

The brute-force approach is to perform a complete in-order traversal at initialization, store all the values in a list, and maintain a pointer to traverse forward and backward. This guarantees correctness because the list will contain all nodes in sorted order, but it requires O(n) space for storage and O(n) time for the initial traversal. Each operation (next and prev) would then take O(1) time.

The optimal approach for the follow-up is to avoid storing the entire traversal upfront. We can use a controlled in-order traversal with a stack to generate elements on demand. This allows next() to run in amortized O(1) time and uses O(h) space for the stack, where h is the tree height. To support prev(), we need to maintain a history of previously visited elements in a list, which allows backward traversal in O(1) time per operation while still limiting space usage to the number of elements already traversed.

Algorithm Walkthrough

1. Initialize the iterator with the BST root. Use a stack to help with on-demand in-order traversal. Traverse the leftmost branch of the tree and push all nodes onto the stack. This prepares the stack so that the smallest element is on top.
2. Maintain a history list to store elements returned by next() or prev(). Keep an index pointer that tracks the current position in the history.
3. next() first checks if the index is pointing to an element already in history. If yes, simply increment index and return the element. If not, pop the top node from the stack, append its value to history, push the leftmost branch of its right child to the stack, increment index, and return the value.
4. prev() decrements index and returns the element from history at the new index. Since we maintain the history list, no tree traversal is needed for backward movement.
5. hasNext() returns true if index is less than the length of history or if the stack is not empty. hasPrev() returns true if index is greater than 0.

Why it works: The stack ensures we always have the next smallest element ready for next(), and the history list guarantees correct prev() operations. The invariant is that all elements before the current index in history are already visited, and the stack always has the next unvisited elements.

Python Solution

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
from typing import Optional

class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.stack = []
        self.history = []
        self.index = -1
        self._push_left_branch(root)

    def _push_left_branch(self, node: Optional[TreeNode]):
        while node:
            self.stack.append(node)
            node = node.left

    def hasNext(self) -> bool:
        return self.index + 1 < len(self.history) or bool(self.stack)

    def next(self) -> int:
        self.index += 1
        if self.index == len(self.history):
            node = self.stack.pop()
            self.history.append(node.val)
            self._push_left_branch(node.right)
        return self.history[self.index]

    def hasPrev(self) -> bool:
        return self.index > 0

    def prev(self) -> int:
        self.index -= 1
        return self.history[self.index]

The Python implementation uses a stack to track the unvisited portion of the BST in in-order sequence. _push_left_branch pushes all left children of a node to the stack. history stores visited elements, allowing prev() to work in constant time. index keeps the position in history.

Go Solution

type BSTIterator struct {
	stack   []*TreeNode
	history []int
	index   int
}

func Constructor(root *TreeNode) BSTIterator {
	it := BSTIterator{stack: []*TreeNode{}, history: []int{}, index: -1}
	it.pushLeftBranch(root)
	return it
}

func (this *BSTIterator) pushLeftBranch(node *TreeNode) {
	for node != nil {
		this.stack = append(this.stack, node)
		node = node.Left
	}
}

func (this *BSTIterator) HasNext() bool {
	return this.index+1 < len(this.history) || len(this.stack) > 0
}

func (this *BSTIterator) Next() int {
	this.index++
	if this.index == len(this.history) {
		n := len(this.stack) - 1
		node := this.stack[n]
		this.stack = this.stack[:n]
		this.history = append(this.history, node.Val)
		this.pushLeftBranch(node.Right)
	}
	return this.history[this.index]
}

func (this *BSTIterator) HasPrev() bool {
	return this.index > 0
}

func (this *BSTIterator) Prev() int {
	this.index--
	return this.history[this.index]
}

In Go, slices are used for both the stack and history. Popping is done by slicing off the last element. Nil checks replace Python's None.

Worked Examples

Consider the tree [7,3,15,null,null,9,20]. The in-order traversal is [3,7,9,15,20].

Complexity Analysis

Amortized time accounts for occasional _push_left_branch traversals when moving to a node’s right subtree.

Test Cases

# test cases
def build_tree(vals):
    if not vals: return None
    nodes = [TreeNode(v) if v is not None else None for v in vals]
    kid_idx = 1
    for i in range(len(nodes)):
        if nodes[i] is not None:
            if kid_idx < len(nodes): nodes[i].left = nodes[kid_idx]; kid_idx += 1
            if kid_idx < len(nodes): nodes[i].right = nodes[kid_idx]; kid_idx += 1
    return nodes[0]

root = build_tree([7,3,15,None,None,9,20])
it = BSTIterator(root)
assert it.next() == 3
assert it.next() == 7
assert it.prev() == 3
assert it.next() == 7
assert it.hasNext() is True
assert it.next() == 9
assert it.next() == 15
assert it.next() == 20
assert it.hasNext() is False
assert it.hasPrev()