Description
设计实现一个带有下列属性的二叉查找树的迭代器:
- 元素按照递增的顺序被访问(比如中序遍历)
next()
和hasNext()
的询问操作要求均摊时间复杂度是O(1)
样例
对于下列二叉查找树,使用迭代器进行中序遍历的结果为 [1, 6, 10, 11, 12]
10
/ \
1 11
\ \
6 12
额外空间复杂度是O(h),其中h是这棵树的高度
Super Star:使用O(1)的额外空间复杂度
Solution
算法思路:
这是一个非常通用的利用 stack 进行 Binary Tree Iterator 的写法。该 Iterator 算法即 non-recursion 的 inorder traversal,不仅仅适用于 BST,任何 Binary Tree 都可以。
- stack 中保存一路走到当前节点的路径中所有的节点。
- stack.peek() 一直指向 iterator 指向的当前节点。
- hasNext() 只需要判断 stack 是否为空。
- next() 只需要返回 stack.peek() 的值,并将 iterator 挪到下一个点,对 stack 进行相应的变化。
挪到下一个点的算法如下:
- 如果当前点存在右子树,那么就是右子树中“一路向西”最左边的那个点。
- 如果当前点不存在右子树,则是走到当前点的路径中,第一个左拐的点。
/**
* Definition of TreeNode:
* public class TreeNode {
* public int val;
* public TreeNode left, right;
* public TreeNode(int val) {
* this.val = val;
* this.left = this.right = null;
* }
* }
* Example of iterate a tree:
* BSTIterator iterator = new BSTIterator(root);
* while (iterator.hasNext()) {
* TreeNode node = iterator.next();
* do something for node
* }
*/
public class BSTIterator {
private Stack<TreeNode> stack = new Stack<>();
/*
* @param root: The root of binary tree.
*/public BSTIterator(TreeNode root) {
// do intialization if necessary
//一直往左边走,找到最左边的节点,并将路径上的点都入栈
while (root != null) {
stack.push(root);
root = root.left;
}
}
/*
* @return: True if there has next node, or false
*/
public boolean hasNext() {
// write your code here
return !stack.isEmpty();
}
/*
* @return: return next node
*/
public TreeNode next() {
// write your code here
TreeNode curt = stack.peek();
TreeNode node = curt;
//调整栈的结构,找到下一个点
//2. 如果当前点不存在右子树,则是走到当前点的路径中,第一个左拐的点
//就是原路返回(stack.pop()),找到第一个左拐的点
//相当于左子树遍历完了,该遍历中间和右边了
if (node.right == null) {
node = stack.pop();
while (!stack.isEmpty() && stack.peek().right == node) {
node = stack.pop();
}
} else {
//1. 如果当前点存在右子树,那么就是右子树中“一路向左”最左边的那个点
//就是先向右走一步,然后一直向左走到头的那个点
node = node.right;
while (node != null) {
stack.push(node);
node = node.left;
}
}
return curt;
}
}