Skip to content

Instantly share code, notes, and snippets.

@qiuwei

qiuwei/TreeIterator

Created Jun 22, 2017
Embed
What would you like to do?
BST TreeIterator
/**
* 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
* }
*/
class Snapshot {
TreeNode treeNode;
boolean toRead;
public Snapshot(TreeNode treeNode, boolean toRead){
this.treeNode = treeNode;
this.toRead = toRead;
}
}
public class BSTIterator {
Stack<Snapshot> stack;
//@param root: The root of binary tree.
public BSTIterator(TreeNode root) {
// write your code here
stack = new Stack<>();
stack.push(new Snapshot(root, false));
}
//@return: True if there has next node, or false
public boolean hasNext() {
// write your code here
return !stack.empty() && stack.peek().treeNode != null;
}
//@return: return next node
public TreeNode next() {
// write your code here
if(stack.empty() || stack.peek().treeNode == null) {
throw new UnsupportedOperationException("No tree node left");
} else {
Snapshot s1 = stack.peek();
if(s1.toRead) {
stack.pop();
return s1.treeNode;
} else {
Snapshot s = stack.peek();
while(!s.toRead) {
//add right ,left and current to stack
stack.pop();
if (s.treeNode.right != null) {
stack.push(new Snapshot(s.treeNode.right, false));
}
stack.push(new Snapshot(s.treeNode, true));
if (s.treeNode.left != null) {
stack.push(new Snapshot(s.treeNode.left, false));
}
s = stack.peek();
}
TreeNode result = stack.peek().treeNode;
stack.pop();
return result;
}
}
}
}
@qiuwei

This comment has been minimized.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment