173. Binary Search Tree Iterator

• categories: Code, Interview Question, leetcode, Medium
• source: https://leetcode.com/problems/binary-search-tree-iterator
• topics: Tree

Description

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

• 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.
• boolean hasNext() Returns true if there exists a number in the traversal to the right of the pointer, otherwise returns false.
• int next() Moves the pointer to the right, then 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() calls will always be valid. That is, there will be at least a next number in the in-order traversal when next() is called.

Ideas

1. Morris Inorder Traversal
2. Using stack
3. Using Python Generator

Code

Morris - fastest, but difficult to understand

class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.current_node = root
 
    def next(self) -> int:
        while self.current_node:
            # If there is no left subtree, visit this node and move to the right subtree.
            if not self.current_node.left:
                next_value = self.current_node.val
                self.current_node = self.current_node.right
                return next_value
            
            # Find the inorder predecessor (rightmost node in the left subtree).
            predecessor = self.current_node.left
            while predecessor.right and predecessor.right != self.current_node:
                predecessor = predecessor.right
 
            if not predecessor.right:
                # Establish a temporary link to the current node.
                predecessor.right = self.current_node
                self.current_node = self.current_node.left
            else:
                # Remove the temporary link and visit the current node.
                predecessor.right = None
                next_value = self.current_node.val
                self.current_node = self.current_node.right
                return next_value
 
    def hasNext(self) -> bool:
        return self.current_node is not None

Stack - easy to understand

class BSTIterator(object):
    def __init__(self, root):
        self.stack = []
        while root:
            self.stack.append(root)
            root = root.left
 
    def hasNext(self):
        return len(self.stack) > 0
 
    def next(self):
        node = self.stack.pop()
        x = node.right
        while x:
            self.stack.append(x)
            x = x.left
        return node.val

Generator - slow but pythonic

class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        self.iter = self._inorder(root)
        self.nxt = next(self.iter, None)
    
    def _inorder(self, node: Optional[TreeNode]) -> Generator[int, None, None]:
        if node:
            yield from self._inorder(node.left)
            yield node.val
            yield from self._inorder(node.right)
 
    def next(self) -> int:
        res, self.nxt = self.nxt, next(self.iter, None)
        return res
 
    def hasNext(self) -> bool:
        return self.nxt is not None