【二叉树】二叉搜索树迭代器

0x00 题目

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

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

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

0x01 思路

可以直接对二叉搜索树做一次完全的递归遍历
获取 ​​​中序遍历​​​ 的全部结果并保存在 ​​数组​​​ 中
随后，我们利用得到的数组本身来实现迭代器

复杂度分析
时间复杂度：初始化需要 ​​​O(n)​​​ 的时间
其中 ​​​n​​​ 为树中节点的数量
随后每次调用只需要 ​​​O(1)​​ 的时间

空间复杂度：​​O(n)​​，因为需要保存中序遍历的全部结果

0x02 解法

语言：​​Swift​​

树节点：​​TreeNode​​

public class TreeNode {
    public var val: Int
    public var left: TreeNode?
    public var right: TreeNode?
    public init() { self.val = 0; self.left = nil; self.right = nil; }
    public init(_ val: Int) { self.val = val; self.left = nil; self.right = nil; }
    public init(_ val: Int, _ left: TreeNode?, _ right: TreeNode?) {
        self.val = val
        self.left = left
        self.right = right
    }
}

解法：

class BSTIterator {
    var index: Int = 0
    var array: [Int] = []

    init(_ root: TreeNode?) {
        inorder(root)
    }

    // 中序遍历
    private func inorder(_ root: TreeNode?) {
        guard let r = root else { return }
        inorder(r.left)
        array.append(r.val)
        inorder(r.right)
    }

    // 下一个元素
    func next() -> Int {
        if index < array.count {
            let res = array[index]
            index += 1
            return res
        }
        return 0
    }

    // 是否还有下一个
    func hasNext() -> Bool {
        return index < array.count
    }
}