Taro Logo

Binary Search Tree Iterator II

Medium
Asked by:
Profile picture
10 views
Topics:
TreesStacks

Implement the BSTIterator class that represents an iterator over the in-order traversal of a binary search tree (BST). Implement the following functions:

  • 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. Therefore, hasPrev() should return false and prev() should return an error.
  • boolean hasNext() Returns true if there exists a number in the traversal to the right of the pointer, or false otherwise.
  • int next() Moves the pointer to the right, returns the number at the pointer, and returns the number at the pointer.
  • boolean hasPrev() Returns true if there exists a number in the traversal to the left of the pointer, or false otherwise.
  • int prev() Moves the pointer to the left, returns the number at the pointer, and returns the number at the pointer.

Notice that by initializing the pointer to a non-existent smallest number, the first call to next() will return the smallest element in the BST.

You may assume that next() and prev() calls will always be valid. That is, there will be at least a next/previous number in the in-order traversal when next()/prev() is called.

Example:

Input
["BSTIterator", "next", "next", "prev", "next", "hasPrev", "prev", "next", "hasNext"]
[[[7, 3, 15, null, null, 9, 20]], [], [], [], [], [], [], [], []]
Output
[null, 3, 7, 3, 7, true, 3, 9, true]

Explanation
BSTIterator bSTIterator = new BSTIterator([7, 3, 15, null, null, 9, 20]);
bSTIterator.next(); // return 3
bSTIterator.next(); // return 7
bSTIterator.prev(); // return 3
bSTIterator.next(); // return 7
bSTIterator.hasPrev(); // return True
bSTIterator.prev(); // return 3
bSTIterator.next(); // return 9
bSTIterator.hasNext(); // return True

Constraints:

  • The number of nodes in the tree is in the range [1, 105].
  • 0 <= Node.val <= 106
  • At most 105 calls will be made to hasNext, next, hasPrev, and prev.

Follow up: Could you solve the problem without precalculating the in-order traversal?

Solution


Clarifying Questions

When you get asked this question in a real-life environment, it will often be ambiguous (especially at FAANG). Make sure to ask these questions in that case:

  1. Can the binary search tree contain duplicate values, and if so, how should they be handled by the iterator?
  2. What range of values can the nodes in the BST hold? Are negative values possible?
  3. If the tree is empty, what should the `hasNext()` and `next()` methods return/do, and what should `hasPrevious()` and `previous()` return/do initially?
  4. Is the input guaranteed to be a valid binary search tree, or do I need to validate that as part of the implementation?
  5. After calling `next()`, then `previous()`, is the iterator expected to return the originally returned node from the `next()` call, and vice versa?

Brute Force Solution

Approach

The brute force approach is about listing every single node in the tree in the order we need them and going through that list. We'll keep this list handy so we can just pick the next one when asked. This is like making a cheat sheet of all the nodes in the right order before the test even starts.

Here's how the algorithm would work step-by-step:

  1. First, take the whole tree and flatten it into a simple ordered list, following the rules of how we normally go through a binary search tree (smallest to largest).
  2. Every time someone asks for the 'next' node, just take the next one from that prepared list.
  3. If they ask for the 'previous' node, go back one node in that same list.
  4. If we reach the end of our list when asking for the 'next' node or the beginning when asking for the 'previous' node, just tell them there isn't one.

Code Implementation

class TreeNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

class BinarySearchTreeIterator:
    def __init__(self, root: TreeNode):
        self.ordered_list = []
        self.current_index = -1
        self.inorder_traversal(root)

    def inorder_traversal(self, root: TreeNode):
        if root:
            self.inorder_traversal(root.left)
            self.ordered_list.append(root.value)
            self.inorder_traversal(root.right)

    def hasNext(self) -> bool:
        return self.current_index < len(self.ordered_list) - 1

    def next(self) -> int:
        # Advance the index to the next element.
        self.current_index += 1
        return self.ordered_list[self.current_index]

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

    def prev(self) -> int:
        # Move the index back to the previous element.
        self.current_index -= 1
        return self.ordered_list[self.current_index]

Big(O) Analysis

Time Complexity
O(n)The initialization flattens the entire binary search tree into a sorted list of n nodes using an in-order traversal. This in-order traversal visits each node exactly once, resulting in O(n) time complexity. The next() and previous() operations simply access elements in this pre-computed list, which takes O(1) time each. Therefore, the dominant factor is the initial flattening of the tree, giving the overall time complexity of O(n).
Space Complexity
O(N)The brute force approach involves flattening the entire binary search tree into a list containing all N nodes. This list is stored in memory to allow for quick access to the next and previous nodes. Therefore, the auxiliary space required is proportional to the number of nodes in the tree, which is N. Consequently, the space complexity is O(N).

Optimal Solution

Approach

The challenge is to navigate a binary search tree and remember the path we've taken so we can efficiently move both forward and backward. We maintain two stacks: one to keep track of the path from the root to the current node, and another to store nodes we've previously visited for backtracking.

Here's how the algorithm would work step-by-step:

  1. To initialize, traverse as far left as possible from the root node, pushing each node onto a stack.
  2. When moving 'next', we return the top node from the stack (the current smallest element), but first, we check if it has a right child.
  3. If it has a right child, we traverse as far left as possible from that right child, pushing those nodes onto the stack. This finds the next smallest element.
  4. Before moving to the next element, we check if we have seen this element already. If so we have to deal with going back.
  5. To move 'previous', we have to move up the stack to previous nodes that are smaller than the current node
  6. Store any visited node into the appropriate stack so you can always go back and forth in the tree efficiently.

Code Implementation

class TreeNode:
    def __init__(self, value):
        self.value = value
        self.left = None
        self.right = None

class BSTIterator:
    def __init__(self, root: TreeNode):
        self.forward_stack = []
        self.backward_stack = []
        self._push_all_left(root)

    def _push_all_left(self, node: TreeNode):
        # Push all left nodes starting from node
        while node:
            self.forward_stack.append(node)
            node = node.left

    def hasNext(self) -> bool:
        return len(self.forward_stack) > 0

    def next(self) -> int:
        node = self.forward_stack.pop()
        self.backward_stack.append(node)

        # If it has a right, push all the way left
        if node.right:
            self._push_all_left(node.right)

        return node.value

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

    def prev(self) -> int:
        node = self.backward_stack.pop()

        # Push current node back to forward stack so next() can work
        self.forward_stack.append(node)

        return node.value

Big(O) Analysis

Time Complexity
O(1) amortizedThe hasNext() and hasPrevious() operations each take O(1) time. The next() operation involves pushing nodes onto the stack during the leftmost traversal of the right subtree and popping a node from the stack. While the leftmost traversal can take O(h) time where h is the height of the tree, each node is visited and pushed/popped at most once. The number of push/pop operations will at most be the number of nodes in the tree which is n. Therefore the time complexity for next() is O(1) amortized. Similarly, the previous() operation amortizes to O(1). The constructor takes O(h) time which is equivalent to O(n) time in the worst case if the tree is skewed. Therefore, each call to next() or previous() will be O(1) amortized.
Space Complexity
O(N)The algorithm uses two stacks: one to store the path from the root to the current node during the initial traversal and subsequent 'next' operations, and another to store previously visited nodes for 'previous' operations. In the worst-case scenario, the first stack could contain all the nodes along the leftmost branch of the tree, which could be of height N in a skewed tree, where N is the number of nodes. Similarly, the second stack could potentially store all the previously visited nodes, also up to N in the worst case, as the tree is traversed back and forth. Therefore, the auxiliary space used by these stacks is proportional to N, resulting in a space complexity of O(N).

Edge Cases

Null root node
How to Handle:
The constructor should handle a null root by initializing an empty state (e.g., empty stack).
Tree with only one node
How to Handle:
hasNext() returns true, next() returns the node's value, and hasNext() then returns false.
Tree with all nodes having the same value
How to Handle:
The in-order traversal logic should still function correctly, visiting all nodes sequentially.
Completely unbalanced tree (skewed left or right)
How to Handle:
Ensure the stack storing previous nodes doesn't exceed memory limits for very deep trees.
Calling next() after hasNext() returns false
How to Handle:
Throw an exception like NoSuchElementException to indicate invalid operation.
Very large tree (millions of nodes)
How to Handle:
Ensure the in-order traversal algorithm scales efficiently and doesn't cause stack overflow in recursion, potentially using an iterative approach.
Calling previous() multiple times when at the smallest value
How to Handle:
Ensure that repeated calls to previous() when at the smallest value correctly handles empty stack edge case or some other marker to prevent errors.
Integer overflow in node values
How to Handle:
Use long to avoid integer overflow when comparing or calculating values if the data type allows for larger range.