173、二叉搜索树迭代器
实现一个二叉搜索树迭代器类BSTIterator ,表示一个按中序遍历二叉搜索树(BST)的迭代器:
- BSTIterator(TreeNode root) 初始化 BSTIterator 类的一个对象。BST 的根节点 root 会作为构造函数的一部分给出。指针应初始化为一个不存在于 BST 中的数字,且该数字小于 BST 中的任何元素。
- boolean hasNext() 如果向指针右侧遍历存在数字,则返回 true ;否则返回 false 。
- int next()将指针向右移动,然后返回指针处的数字。
注意,指针初始化为一个不存在于 BST 中的数字,所以对 next() 的首次调用将返回 BST 中的最小元素。
你可以假设 next() 调用总是有效的,也就是说,当调用 next() 时,BST 的中序遍历中至少存在一个下一个数字。
示例1:
输入
["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]
思路:
一看到迭代器,就想到了python的迭代器,那就实现一下吧。
class BSTIterator:
def __init__(self, root: TreeNode):
self.root = root
self.nextVal = None
# 生成迭代对象
self.it = iter(self)
# 实现iter方法,作为其迭代器
def __iter__(self):
cur = self.root
stack = []
while cur or stack:
while cur:
stack.append(cur)
cur = cur.left
cur = stack.pop()
yield cur.val
cur = cur.right
# 由于python迭代器没有hasNext方法,所以这里绕一下,next的Stop判断有无下一个元素
# 这里就需要记录当前迭代的值,如果当前有值,说明已经调用过hasNext,将该值返回即可
# 并且将缓存值设置为None
def next(self) -> int:
if self.nextVal == None: self.hasNext()
nextVal = self.nextVal
self.nextVal = None
return nextVal
# 如果有值说明之前已经调用过hasNext,那么返回True即可
# 如果没有根据next产生的异常,来判断是否有下一个元素
# 这里为什么不直接用stack的长度来判断,由于yield之前将栈元素弹出了,等到再次调用next才会进行下一次循环
# 所以可能stack为空,但是树只是遍历到根节点而已
def hasNext(self) -> bool:
if self.nextVal != None: return True
try:
self.nextVal = next(self.it)
return True
except StopIteration:
return False
这样看起来,用原生的迭代器实现还是比较绕的,那么用迭代方式实现吧!
class BSTIterator:
# 初始化,遍历到左叶子节点
def __init__(self, root: TreeNode):
self.stack = []
cur = root
while cur:
self.stack.append(cur)
cur = cur.left
def next(self) -> int:
# cur即是中序遍历的结点,弹出后,需要进行右子节点的遍历
cur = self.stack.pop()
node = cur.right
while node:
self.stack.append(cur)
cur = cur.left
return cur.val
# 由于每次next都会去找到右子结点的左节点,所以stack不为空
# 如果为空,说明树已经遍历完成
def hasNext(self) -> bool:
return len(self.stack) > 0