Binary Search Tree Iterator
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()
andhasNext()
queries run in O(1) time in average.Example
For the following binary search tree, in-order traversal by using iterator is [1,
6, 10, 11, 12]
10
/ \
1 11
\ \
6 12
Extra memory usage O(h), h is the height of the tree.
Super Star: Extra memory usage O(1)
/**
* 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:
* Solution iterator = new Solution(root);
* while (iterator.hasNext()) {
* TreeNode node = iterator.next();
* do something for node
* }
*/
public class Solution {
public Stack myStack = new Stack();
public TreeNode current;
//@param root: The root of binary tree.
public Solution(TreeNode root) {
// write your code here
current = root;
return;
}
//@return: True if there has next node, or false
public boolean hasNext() {
// write your code here
return (current!=null || !myStack.isEmpty());
}
//@return: return next node
public TreeNode next() {
// write your code here
while(current!=null){
myStack.push(current);
current = current.left;
}
TreeNode res = (TreeNode) myStack.pop();
current = res.right;
return res;
}
}
核心:TreeNode next()
主要功能一直找出最小的那个node:
1. 如果current != null,一直找left.left.left.....就是所求结果;
2. 如果current == null 并且current是 current.parent的left node时,current.parent 就是所求结果,也就是myStack最上面的node;
3. 如果current == null 并且current是 current.parent的right node时, current.parent已经被pop了,current.parent.parent是所求结果,也是myStack最上面的node。
第一个while loop非常巧妙:当current不等于null时,一直查询left node,并把沿途left node存入myStack,一直到最小的node的left node (值为null),然后再从stack里面读出来上一个left node,虽然可以用while(current.left!=null)来省略两步(current=current.left 和 current=myStack.pop() ), 但这样的话current == null的情况要多写一个if语句。while(curerent != null) 包括了 if(current == null) 的情况。 之后的从myStack读出来是两个情况都要做的下一步。
current 为结果时,再从以 current.right 为 root 的 tree 中找 next(), 所以最后要 current = current.right