难度: Medium
Submission
运行时间: 295 ms
内存: 48.9 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 = stack = []
while root:
stack.append(root)
root = root.left
self.node = None
def hasNext(self) -> bool:
if not self.node: return bool(self.stack)
if self.node.right: return True
node = self.node
for p in reversed(self.stack):
if p.left == node: return True
node = p
return False
def next(self) -> int:
stack = self.stack
if not self.node:
self.node = stack.pop()
elif self.node.right:
stack.append(self.node)
node = self.node.right
while node.left:
stack.append(node)
node = node.left
self.node = node
else:
if not stack: return self.node.val
node = self.node
while stack[-1].right == node:
node = stack.pop()
self.node = stack.pop()
return self.node.val
def hasPrev(self) -> bool:
if not self.node: return False
if self.node.left: return True
node = self.node
for p in reversed(self.stack):
if p.right == node: return True
node = p
return False
def prev(self) -> int:
node = self.node
stack = self.stack
if node.left:
stack.append(node)
node = node.left
while node.right:
stack.append(node)
node = node.right
self.node = node
elif not self.hasPrev():
self.stack.append(node)
self.node = None
return node.val
else:
while stack[-1].left == node:
node = stack.pop()
self.node = stack.pop()
return self.node.val
# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.hasNext()
# param_2 = obj.next()
# param_3 = obj.hasPrev()
# param_4 = obj.prev()
Explain
这个题解实现了一个二叉搜索树的双向迭代器，支持获取当前节点的前驱和后继节点。主要思路是： 1. 初始化时，将根节点到最左叶子节点的路径压入栈中。 2. 获取下一个节点时，如果当前节点有右子树，则将当前节点压栈，并将右子树的最左路径压入栈，当前节点指向右子树的最左叶子；否则弹出栈顶节点，直到栈顶节点是当前节点的父节点，当前节点指向该父节点。 3. 获取上一个节点时，如果当前节点有左子树，则把当前节点压栈，并将左子树的最右路径压入栈，当前节点指向左子树的最右叶子；否则弹出栈顶节点，直到栈顶节点是当前节点的父节点，当前节点指向该父节点。 4. 判断是否有下一个节点，当前节点为空时，看栈是否为空；当前节点不为空时，若有右子树则必有下一个，否则看栈中节点是否为当前节点的父节点。 5. 判断是否有上一个节点，若当前节点为空则没有；若当前节点有左子树则必有上一个，否则看栈中节点是否为当前节点的父节点。
时间复杂度: 初始化：O(n)；next()和prev()：平均O(1)，最坏O(n)；hasNext()和hasPrev()：平均O(1)，最坏O(n)
空间复杂度: 平均空间复杂度是O(log(n))，最坏情况是O(n)
class BSTIterator:
def __init__(self, root: Optional[TreeNode]):
self.stack = stack = []
# 初始化时将根到最左叶子的路径节点压入栈
while root:
stack.append(root)
root = root.left
# 当前节点初始为空
self.node = None
def hasNext(self) -> bool:
# 若当前节点为空，根据栈是否为空判断是否有后继
if not self.node: return bool(self.stack)
# 若当前节点有右子树，则必然有后继节点
if self.node.right: return True
node = self.node
# 检查栈中节点是否为当前节点的父节点
for p in reversed(self.stack):
if p.left == node: return True
node = p
return False
def next(self) -> int:
stack = self.stack
# 若当前节点为空，则后继为栈顶节点
if not self.node:
self.node = stack.pop()
# 若当前节点有右子树，则后继在右子树中
elif self.node.right:
stack.append(self.node)
node = self.node.right
# 进入右子树的最左路径
while node.left:
stack.append(node)
node = node.left
self.node = node
else:
if not stack: return self.node.val
node = self.node
# 若栈中节点是当前节点右孩子，弹出
while stack[-1].right == node:
node = stack.pop()
# 后继为栈顶节点
self.node = stack.pop()
return self.node.val
def hasPrev(self) -> bool:
# 若当前节点为空，则没有前驱
if not self.node: return False
# 若当前节点有左子树，则必有前驱
if self.node.left: return True
node = self.node
# 检查栈中节点是否为当前节点的父节点
for p in reversed(self.stack):
if p.right == node: return True
node = p
return False
def prev(self) -> int:
node = self.node
stack = self.stack
# 若当前节点有左子树，则前驱在左子树中
if node.left:
stack.append(node)
node = node.left
# 进入左子树的最右路径
while node.right:
stack.append(node)
node = node.right
self.node = node
# 若没有前驱节点，将当前节点压栈，当前节点置空
elif not self.hasPrev():
self.stack.append(node)
self.node = None
return node.val
else:
# 若栈顶节点是当前节点左孩子，弹出
while stack[-1].left == node:
node = stack.pop()
# 前驱节点为栈顶元素
self.node = stack.pop()
return self.node.val
Explore
在初始化阶段，迭代器通过从根节点开始，一直向左遍历到最左叶子节点的过程中，将遍历的每个节点压入栈中。这种方法确保了迭代器能够处理包括空树和只有一个节点的树在内的各种情形。对于空树，初始化过程中根节点为`None`，因此栈将保持空状态，表示迭代器没有更多元素。对于只有一个节点的树，该节点将被压入栈中，成为迭代器的起始点。这样的初始化确保了迭代器的正确性和鲁棒性。
在`next`方法中，如果当前节点为空且栈也为空，这表示迭代器已经遍历完所有元素，没有更多的后继节点可以返回。在这种情况下，按照常见的编程实践，应当抛出一个异常（例如`StopIteration`），提醒使用者所有元素已经被遍历完毕。这与Python的迭代器协议一致，当迭代器中没有更多元素时，应当通过抛出异常来通知调用者。
在`prev`方法中，如果没有前驱节点，选择将当前节点压栈并将当前节点置空是为了保持迭代器的状态一致性。这样做可以确保当再次调用`next`方法时，可以从栈中恢复之前的状态，继续正确地进行迭代。此外，将当前节点置空也是为了防止重复访问同一个节点，确保每次调用`prev`方法都能正确地返回前一个元素或正确地表达没有更多的前驱节点。这种处理方式有助于迭代器在前后移动时保持其行为的一致性和预测性。
在`hasNext`和`hasPrev`方法中，判断栈中的节点是否为当前节点的父节点可以通过回溯栈的方式实现。具体来说，可以从栈顶开始向栈底回溯，比较每个节点与当前节点的关系。在`hasNext`方法中，如果栈中的某个节点是当前节点的左孩子，这表示该栈中节点是当前节点的父节点，因此存在后继节点。类似地，在`hasPrev`方法中，如果栈中的节点是当前节点的右孩子，这同样表明该栈中节点是当前节点的父节点，因此存在前驱节点。这种基于栈的回溯方法有效地解决了在复杂树结构中确定父子关系的问题，保证了迭代器逻辑的正确性。