python-leetcode-二叉搜索树迭代器

173. 二叉搜索树迭代器 - 力扣（LeetCode）

# 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
class BSTIterator:
    def __init__(self, root: TreeNode):
        # 初始化时用栈来模拟中序遍历
        self.stack = []
        # 将左子树的所有节点压入栈
        self._push_all_left_nodes(root)
    
    # 辅助函数：将所有的左子节点压入栈
    def _push_all_left_nodes(self, node):
        while node:
            self.stack.append(node)
            node = node.left

    def hasNext(self) -> bool:
        # 判断栈是否为空
        return len(self.stack) > 0

    def next(self) -> int:
        # 弹出栈顶元素
        node = self.stack.pop()
        # 如果弹出的节点有右子树，处理右子树
        if node.right:
            self._push_all_left_nodes(node.right)
        # 返回当前节点的值
        return node.val
        


# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()