leetcode树之二叉搜索树迭代器_leetcode二叉搜索树迭代器

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]
1
2
3
4
5

---

思路：

一看到迭代器，就想到了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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

这样看起来，用原生的迭代器实现还是比较绕的，那么用迭代方式实现吧！

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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21