Description
Design an iterator that implements a binary search tree with the following properties:
- Elements are accessed in increasing order (such as in-order traversal)
next()
hasNext()
The query operation of sum requires amortized time complexity of O( 1 )
Sample
For the following binary search tree, the result of an inorder traversal using an iterator is [1, 6, 10, 11, 12]
10
/ \
1 11
\ \
6 12
challenge
The extra space complexity is O(h), where h is the height of the tree
Super Star: extra space complexity using O(1)
Solution
Algorithm idea:
This is a very general use of stack to write Binary Tree Iterator. The Iterator algorithm is non-recursion inorder traversal, not only applicable to BST, but any Binary Tree.
- The stack stores all the nodes in the path all the way to the current node.
- stack.peek() always points to the current node pointed to by iterator.
- hasNext() only needs to judge whether the stack is empty.
- next() only needs to return the value of stack.peek(), move the iterator to the next point, and make corresponding changes to the stack.
The algorithm for moving to the next point is as follows:
- If the current point has a right subtree, then it is the leftmost point in the right subtree "all the way west".
- If there is no right subtree at the current point, it is the first left-turn point in the path to the current point.
/** * 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 //Go all the way to the left, find the leftmost node, and push the points on the path to the stack 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; //Adjust the structure of the stack to find the next point //2. If there is no right subtree at the current point, it is the first left turn point in the path to the current point //It is to return to the original road (stack.pop()) and find the first left turn point // Equivalent to the left subtree traversed, it is time to traverse the middle and the right if (node.right == null) { node = stack.pop(); while (!stack.isEmpty() && stack.peek().right == node) { node = stack.pop(); } } else { //1. If there is a right subtree at the current point, then it is the leftmost point "all the way to the left" in the right subtree //It is to take one step to the right, and then go all the way to the left to the point of the head node = node.right; while (node != null) { stack.push(node); node = node.left; } } return curt; } }