二叉搜索树迭代器

标签: 栈 树 设计 二叉搜索树 二叉树 迭代器

难度: Medium

实现一个二叉搜索树迭代器类BSTIterator ，表示一个按中序遍历二叉搜索树（BST）的迭代器：
  • BSTIterator(TreeNode root) 初始化 BSTIterator 类的一个对象。BST 的根节点 root 会作为构造函数的一部分给出。指针应初始化为一个不存在于 BST 中的数字，且该数字小于 BST 中的任何元素。
  • boolean hasNext() 如果向指针右侧遍历存在数字，则返回 true ；否则返回 false 。
  • int next()将指针向右移动，然后返回指针处的数字。

注意，指针初始化为一个不存在于 BST 中的数字，所以对 next() 的首次调用将返回 BST 中的最小元素。

你可以假设 next() 调用总是有效的，也就是说，当调用 next() 时，BST 的中序遍历中至少存在一个下一个数字。

 

示例：

输入
["BSTIterator", "next", "next", "hasNext", "next", "hasNext", "next", "hasNext", "next", "hasNext"]
[[[7, 3, 15, null, null, 9, 20]], [], [], [], [], [], [], [], [], []]
输出
[null, 3, 7, true, 9, true, 15, true, 20, false]

解释
BSTIterator bSTIterator = new BSTIterator([7, 3, 15, null, null, 9, 20]);
bSTIterator.next();    // 返回 3
bSTIterator.next();    // 返回 7
bSTIterator.hasNext(); // 返回 True
bSTIterator.next();    // 返回 9
bSTIterator.hasNext(); // 返回 True
bSTIterator.next();    // 返回 15
bSTIterator.hasNext(); // 返回 True
bSTIterator.next();    // 返回 20
bSTIterator.hasNext(); // 返回 False

 

提示：

  • 树中节点的数目在范围 [1, 105] 内
  • 0 <= Node.val <= 106
  • 最多调用 105 次 hasNext 和 next 操作

 

进阶：

  • 你可以设计一个满足下述条件的解决方案吗？next() 和 hasNext() 操作均摊时间复杂度为 O(1) ，并使用 O(h) 内存。其中 h 是树的高度。

Submission

运行时间: 80 ms

内存: 21.8 MB

# 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: Optional[TreeNode]):
        self.stack = []
        self.dfsLeft(root)

    def next(self) -> int:
        root = self.stack.pop()
        self.dfsLeft(root.right)
        return root.val

    def hasNext(self) -> bool:
        return len(self.stack) != 0

    def dfsLeft(self, root):
        while root is not None:
            self.stack.append(root)
            root = root.left


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

Explain

这个题解使用了迭代的方式实现二叉搜索树的中序遍历。在初始化时，通过 dfsLeft 方法将根节点及其所有左子节点压入栈中。在调用 next 方法时，从栈中弹出一个节点，该节点就是当前中序遍历的最小节点。然后对该节点的右子树调用 dfsLeft，将右子树的所有左子节点压入栈中。hasNext 方法通过判断栈是否为空来确定是否还有下一个元素。

时间复杂度: 平均情况下 O(1)，最坏情况下 O(n)

空间复杂度: 平均情况下 O(h)，最坏情况下 O(n)

# 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: Optional[TreeNode]):
        self.stack = []
        self.dfsLeft(root)  # 初始化时将所有左子节点压入栈

    def next(self) -> int:
        root = self.stack.pop()  # 弹出栈顶元素，即当前最小节点
        self.dfsLeft(root.right)  # 对弹出节点的右子树调用 dfsLeft
        return root.val

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

    def dfsLeft(self, root):
        while root is not None:  # 将一个节点的所有左子节点压入栈
            self.stack.append(root)
            root = root.left


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

Explore

在二叉搜索树中，中序遍历会按照从小到大的顺序访问所有节点。将所有左子节点先压入栈中是为了保证每次调用`next()`方法时能够按照这种顺序返回最小的元素。由于栈是后进先出的数据结构，通过先将左子节点压栈，确保了每次从栈中弹出的是当前未访问节点中值最小的节点。这种方式简化了中序遍历的复杂度，使得每次调用`next()`的时间复杂度为O(1)（平均情况下）。

在`dfsLeft`方法中，当遇到一个节点没有左子树时，该节点本身会被压入栈中，然后方法终止对该节点的左子树的进一步遍历，因为左子树为空。随后，`next()`方法会将这个节点从栈中弹出，并根据需要对其右子树应用同样的`dfsLeft`处理，即尝试找到右子树中所有的左子节点并压栈。此处理确保了即使节点没有左子树，其本身和右子树中的节点也能被正确访问和返回。

在中序遍历中，一个节点的访问顺序是先左子树，然后节点本身，最后右子树。当`next()`方法从栈中弹出一个节点时，表示该节点的左子树已经完全访问完毕，接下来需要访问该节点本身（这已经通过`next()`实现），然后是其右子树。为了维持中序遍历的顺序，我们需要立即对弹出节点的右子树执行`dfsLeft`，确保其右子树中所有左侧节点被正确压入栈中，从而保证下一次调用`next()`时能够访问到正确的最小节点。

是的，这种实现方式即便在二叉搜索树中存在重复值时也能正确运行。在二叉搜索树中，即使值重复，每个节点仍然会保持固定的位置，使得中序遍历的顺序依然是定义良好的。在实际实现中，只要树的结构正确维护（即使存在值的重复），迭代器将按照中序遍历的逻辑返回每个节点，包括那些值重复的节点。