题目
Design an iterator over a binary search tree with the following rules:
Elements are visited in ascending order (i.e. an in-order traversal)
next() and hasNext() queries run in O(1) time in average.
我的想法
看懂题是关键。
用一个stack记录目前经过了的点,并且栈头为目前的最小元素,可以控制空间复杂度为树的高度。
找next只需把当前的栈顶pop出来,再把当前栈顶的右子树的所有左子树加入栈即可
/**
* 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 {
Stack<TreeNode> stack;
/*
* @param root: The root of binary tree.
*/public BSTIterator(TreeNode root) {
stack = new Stack<>();
while(root != null) {
stack.push(root);
root = root.left;
}
}
/*
* @return: True if there has next node, or false
*/
public boolean hasNext() {
if(stack.isEmpty()) {
return false;
}
return true;
}
/*
* @return: return next node
*/
public TreeNode next() {
if(!hasNext()) {
return null;
}
TreeNode next = stack.pop();
TreeNode node = next.right;
while(node != null) {
stack.push(node);
node = node.left;
}
return next;
}
}