实现一个二叉搜索树迭代器类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
解法一
先不考虑题目 Note
中要求的空间复杂度和时间复杂度,简单粗暴一些。在构造函数中,对二叉树进行中序遍历,把结果保存到一个队列中,然后 next
方法直接执行出队操作即可。至于 hasNext
方法的话,判断队列是否为空即可。
class BSTIterator {
Queue<Integer> queue = new LinkedList<>();
public BSTIterator(TreeNode root) {
inorderTraversal(root);
}
private void inorderTraversal(TreeNode root) {
if (root == null) {
return;
}
inorderTraversal(root.left);
queue.offer(root.val);
inorderTraversal(root.right);
}
/** @return the next smallest number */
public int next() {
return queue.poll();
}
/** @return whether we have a next smallest number */
public boolean hasNext() {
return !queue.isEmpty();
}
}
时间复杂度的话,构造函数因为遍历了一遍二叉树,所以是 O(n)
,对于 next
和 hasNext
方法都是 O(1)
。
空间复杂度,用队列保存了所有的节点值,所以是 O(n)
,此外中序遍历递归压栈的过程也需要 O(h)
的空间。
解法二
解法一中我们把所有节点都保存了起来,其实没必要一次性保存所有节点,而是需要一个输出一个即可。
所以我们要控制中序遍历的进程,不要让它一次性结束,如果用解法一递归的方法去遍历那就很难控制了,所以自然而然的会想到用栈模拟递归的过程。
下边是 94 题 解法二的代码。
public List<Integer> inorderTraversal(TreeNode root) {
List<Integer> ans = new ArrayList<>();
Stack<TreeNode> stack = new Stack<>();
TreeNode cur = root;
while (cur != null || !stack.isEmpty()) {
//节点不为空一直压栈
while (cur != null) {
stack.push(cur);
cur = cur.left; //考虑左子树
}
//节点为空,就出栈
cur = stack.pop();
//当前值加入
ans.add(cur.val);
//考虑右子树
cur = cur.right;
}
return ans;
}
和这道题糅合一起也很简单了,只需要把 stack
和 cur
作为成员变量,然后每次调用 next
就执行一次 while
循环,并且要记录当前值,结束掉本次循环。
class BSTIterator {
Stack<TreeNode> stack = new Stack<>();
TreeNode cur = null;
public BSTIterator(TreeNode root) {
cur = root;
}
/** @return the next smallest number */
public int next() {
int res = -1;
while (cur != null || !stack.isEmpty()) {
// 节点不为空一直压栈
while (cur != null) {
stack.push(cur);
cur = cur.left; // 考虑左子树
}
// 节点为空,就出栈
cur = stack.pop();
res = cur.val;
// 考虑右子树
cur = cur.right;
break;
}
return res;
}
/** @return whether we have a next smallest number */
public boolean hasNext() {
return cur != null || !stack.isEmpty();
}
}
时间复杂度的话,对于 next
方法,大多数时候是 O(1)
,但最坏情况因为最里边的 while
循环,其实有可能达到 O(n)
。但如果算均摊时间复杂度的话,其实还是 O(1)
,因为每个节点最多也就经过两次就出栈了。
空间复杂度,这里只需要消耗栈的空间,也就是 O(h)
。