173. 二叉搜索树迭代器
实现一个二叉搜索树迭代器类 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 是树的高度。
方法一:保存全部元素
解题思路
- 一种容易想到的思路是对树进行「中序遍历」,使用队列把元素都提前保存起来,next 方法移动指针即可。
参考代码
class BSTIterator {
private int idx;
private final List<Integer> queue;
public BSTIterator(TreeNode root) {
idx = 0;
queue = new ArrayList<>();
inorder(root);
}
public int next() {
return queue.get(idx++);
}
public boolean hasNext() {
return idx < queue.size();
}
// 二叉树中序遍历
private void inorder(TreeNode root) {
Stack<TreeNode> stack = new Stack<>();
while (root != null || !stack.isEmpty()) {
while (root != null) {
stack.push(root);
root = root.left;
}
root = stack.pop();
queue.add(root.val);
root = root.right;
}
}
}
执行结果
- 时间复杂度:BSTIterator - O(n);next - O(1);hasNext - O(1)
- 空间复杂度:O(n)
方法二:只保存左节点
思路
- 在方法一中,我们提前的存储了所有节点,导致空间复杂度为 O(n)。事实上,根据中序遍历的特点:左 -> 中 -> 右,只需要存储左节点即可。
参考代码
class BSTIterator {
private final Deque<TreeNode> stack;
public BSTIterator(TreeNode root) {
stack = new LinkedList<>();
pushLeft(root);
}
public int next() {
TreeNode node = stack.pop();
int ret = node.val;
pushLeft(node.right);
return ret;
}
public boolean hasNext() {
return !stack.isEmpty();
}
private void pushLeft(TreeNode root) {
while (root != null) {
stack.push(root);
root = root.left;
}
}
}
执行结果
- 时间复杂度:next - O(1);hasNext - O(1)。next 方法中,平均每个元素只进出栈一次,均摊复杂度为 O(1)。
- 空间复杂度:O(k),最多只需保存树高度个数的节点。