1586. Binary Search Tree Iterator II

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.
boolean hasPrev() Returns true if there exists a number in the traversal to the left of the pointer, otherwise returns false.
int prev() Moves the pointer to the left, 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() 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.

Idea

用栈实现BST的遍历
同时用两个栈：

• 一个记录接下来要遍历的节点 – stack
• 一个记录已经遍历过的节点 – next_stack

同时initialize一个index， 指向当前节点的位置。每次从next_stack 里pop出一个节点时， 都需要将这个index reset成 指向 stack的最后一个节点。通过index的位置去判断hasNext和hasPrev。

• 如果next_stack有值或当前index不是指向stack的最后一个节点 则 hasNext = True
• 如果index不是0, index >= 1则说明在stack里有当前节点的前继节点， hasPrev = False

Code

class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        
        self.next_stack = [] # 记录所有还没有被 visited 的node
        self.stack = [] # 记录所有被 visited 的node
        self.ind = 0 #当前node的位置
        # 从根节点开始把所有左边节点放进stack
        # 从大到小放到nextstack， nextstack的最后一个是当前最小节点
        while root:
            self.next_stack.append(root)
            root = root.left

    def hasNext(self) -> bool:
        # 还有节点没有被visited OR 所有节点都被visited 但是当前节点不是最后一个节点
        return len(self.next_stack) > 0 or self.ind + 1 < len(self.stack)
        

    def next(self) -> int:
        
        # 当前节点是已经在stack里 且不是最后一个
        if self.ind + 1 < len(self.stack):
            self.ind += 1
            return self.stack[self.ind].val
        
        # 当前节点是stack里的最后一个， 继续看nextstack里是否有下一个节点
        node = self.next_stack.pop()
        
        # 把nextstack里的节点pop出来加到stack里
        # 把index reset指向stack里的最后一个节点 -- 当前节点
        self.stack.append(node)
        self.ind = len(self.stack) - 1
        
        # 如果当前node有右子树，把右子树的左节点加到next里，从大到小的顺序从而使 nextstack里的最新一个node是当前最小的
        current = node.right
        while current:
            self.next_stack.append(current)
            current = current.left
        
        return node.val
        
    def hasPrev(self) -> bool:
        # 如果index >= 1 则说明stack里有前一个节点 stack[0]
        return self.ind != 0
        
    def prev(self) -> int:
        # 把当前index往前移1 然后return
        self.ind -= 1
        return self.stack[self.ind].val

Notes

如果是把BST从大到小的顺序pop出来， 则把所有.right 和 .left 对换即可。